[michakraus]

Summary

This PR adds the parametric core of the generic layers of #207 (relates to #207): apply_parametric, ParametricLoss and ParametricDataLoader. The three names are public and not exported.

• apply_parametric(model, x, μ, ps) applies a Chain to one sample (a vector or a (q, p) pair of vectors) with one μ, or to a batch (a 2d × 1 × k tensor or a (q, p) pair of d × 1 × k tensors) with a vector of k values of μ. The output has the shape of the input. A batch maps over the columns of its d × k reshape through the internal map_samples, so the backward pass allocates O(k), not O(k²). map_samples has a ChainRulesCore.rrule that sums the sample gradients of ps in one NetworkParameters, so the gradient of ParametricLoss on an array batch infers a concrete type. The rule gives the mapped closure NoTangent(): a batch has no gradient with respect to the fields of model, and ZygotePullback differentiates ps only. The internal concatenate_array_with_parameters puts the flattened μ below the state, in the element type of the state.
• ParametricLoss() is the relative norm of the error of apply_parametric. ZygotePullback takes the triple (input, output, μ) with one μ or a vector of them.
• ParametricDataLoader takes a 2d × nt × m array, a (q, p) pair of d × nt × m arrays of one array type, or an EnsembleSolution, with one named tuple of system parameters per trajectory. It batches on main's gather and batch_over_two_axes. train! on it trains on ParametricLoss() when no loss is given.
• The loader raises an ArgumentError that names the numbers for: a Batch with a sequence length or a prediction window other than one, a parameter vector of the wrong length, no trajectory, too few time steps, a μ of another element type, and a μ of another shape. A (q, p) pair of two array types is a MethodError, as for DataLoader (K20 in KNOWN_ISSUES.md).

The code is ported from #207 by Benedikt Brantner, with the fix of its bug 1 (hcat, not vcat) kept and tested. Differences from #207: concatenate_array_with_parameters keeps a Float32 state in Float32, repeats one μ in every column of a matrix, and the counts are checked with an ArgumentError instead of @assert. The batching of the 2d × nt × m array form is new.

Tests

New test files: test/parametric.jl, test/loss/parametric_loss.jl, test/data_loader/parametric_data_loader.jl, with the helper test/helpers/parametric_affine.jl. Every expected value is a closed form built by hand in the test. The gradients are checked against central finite differences, with respect to the network parameters and to the input. All numeric tests run in Float32 and Float64. test/parametric.jl also holds an allocation-difference test of the backward pass across 64 and 1024 samples, for the array form and the (q, p) form.

Known limitations

• The matrix–vector method of concatenate_array_with_parameters (src/parametric.jl:28-33) indexes one column per sample, and its pullback allocates O(k²). apply_parametric does not call it: it passes one vector per sample.
• Three of the four _flatten_system_parameters methods have no caller in this PR. The function is ported from Parametric Generalized Hamiltonian Neural Networks #207 as one unit, and its callers arrive with the later parts of the split of Parametric Generalized Hamiltonian Neural Networks #207.
• A μ with a non-numeric leaf raises the error of NeuralNetworkParameters, not a message of the loader.
• The loss norm of AbstractNeuralNetworks promotes a Float32 (q, p) named tuple to Float64 (upstream, K28 in KNOWN_ISSUES.md).
• The gradient of ParametricLoss on a (q, p) batch still infers as Union{Nothing, NetworkParameters}. The cause is _compute_loss in AbstractNeuralNetworks, and main shows the same (upstream, K30 in KNOWN_ISSUES.md).

🤖 Generated with Claude Code

[coderabbitai]

📝 Summary

Summary by CodeRabbit

• New Features

  • Added support for applying networks with system parameters supplied separately, for single inputs and batches.
  • Added parameter-aware loading of trajectory and ensemble data, pairing each sample with its trajectory’s parameters.
  • Added a relative L₂ loss for parametric training. Training with the new loader uses this loss by default when none is specified.
  • Added pullback support for losses that include system parameters.

• Documentation

  • Added guidance on parametric application, loss, data loading, accepted inputs, and training.

• Known Issues

  • Documented inference limitations for certain batched inputs and a precision-promotion issue affecting some Float32 data.

This change adds parametric model application, a trajectory loader that carries system parameters, a relative L₂ loss, and pullback support. Training accepts the new loader and uses ParametricLoss() by default. Tests and documentation cover the APIs, input shapes, and observed known issues.

Changes: Parametric system-parameter training

Layer / File(s)	Summary
Parametric model application src/parametric.jl, src/GeometricMachineLearning.jl, test/parametric.jl, test/helpers/parametric_affine.jl	Adds helpers to flatten and append system parameters, plus single-sample and batch apply_parametric methods. Tests cover array and (q, p) inputs, shape and type behavior, gradients, and allocation scaling.
Parametric trajectory loading src/data_loader/parametric_data_loader.jl, src/GeometricMachineLearning.jl, test/data_loader/parametric_data_loader.jl, docs/src/data_loader/data_loader.md	Adds ParametricDataLoader constructors for array, (q, p), and ensemble data. It validates trajectory and parameter shapes, creates one-step batches with states, next-step targets, and parameters, and adds loader tests and documentation.
Parametric loss and pullbacks src/loss/losses.jl, src/pullbacks/zygote_pullback.jl, test/loss/parametric_loss.jl, test/parametric.jl, test/runtests.jl, docs/src/reduced_order_modeling/losses.md, KNOWN_ISSUES.md	Adds ParametricLoss and pullback methods for one system-parameter value or a vector of per-sample values. Tests check loss values, gradients, and element types. The known-issues entries record inference and Float32 promotion observations.
Training integration src/data_loader/train.jl, test/data_loader/parametric_data_loader.jl, test/runtests.jl, CHANGELOG.md	train! accepts ParametricDataLoader and selects ParametricLoss() as its default loss. Tests compare default training with an explicit ParametricLoss(), and the changelog describes the API and accepted data shapes.

Priority: Normal

Estimated code review effort: 3 (Moderate) | ~25 minutes

Change: Feature

Sequence Diagram(s)

sequenceDiagram
  participant Training as train!
  participant Loader as ParametricDataLoader
  participant Loss as ParametricLoss
  participant Apply as apply_parametric
  participant Model as Chain
  Training->>Loader: Obtain batches with states, targets, and parameters
  Training->>Loss: Evaluate the model on a batch
  Loss->>Apply: Apply model with input and system parameters
  Apply->>Model: Forward pass with state, system parameters, and network parameters

Loading

Possibly related issues

• Take over the parametric dataset helpers from GeometricMachineLearning's scripts JuliaGML/GMLDatasets.jl#5: The issue describes a parametric dataset interface and identifies ParametricDataLoader as a target for shared parameter-formatting and (q, p) reshaping helpers.
• No stated or tested contract that a gradient leaf lands on its parameter's backend AbstractNeuralNetworks.jl#39: The issue describes a backend-placement contract for ZygotePullback parameter gradients, which is relevant to the added parametric pullback methods.

Possibly related PRs

• JuliaGNI/GeometricMachineLearning.jl#353: This PR uses NamedTupleOfArrays and ArrayOrNamedTuple types introduced by #353 for parametric application, loading, loss, and pullback signatures.
• JuliaGNI/GeometricMachineLearning.jl#356: Both PRs document the Float32 (q, p) loss-norm promotion behavior in KNOWN_ISSUES.md.

Merge Risk: Low · up to c51e9

Parametric training has narrow edge-case failures for empty batches and supported custom non-Chain models; both have localized fixes and explicit-loss workarounds exist for the latter.

Pre-merge checks: 5

Check name	Status	Explanation
Title check		The title names the three public entities added by the PR (apply_parametric, ParametricLoss, ParametricDataLoader) and the main change, the parametric core. It matches the changeset and the PR objecti…
Docstring Coverage		No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check. Docstring coverage is scoped to functions touched by this diff. Analyzed 0 functions across 0…
Linked Issues check		Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check		Check skipped because no linked issues were found for this pull request.
Description check		The pull request description directly matches the changeset. It describes apply_parametric, ParametricLoss, ParametricDataLoader, train! integration, tests, and known limitations.

Full details: Title check

Explanation

The title names the three public entities added by the PR (apply_parametric, ParametricLoss, ParametricDataLoader) and the main change, the parametric core. It matches the changeset and the PR objectives. It is 79 characters and fits a concise, specific title.

Full details: Docstring Coverage

Explanation

No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check. Docstring coverage is scoped to functions touched by this diff. Analyzed 0 functions across 0 files. (2 skipped: 2 unsupported.)

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 1

---

Review info

Run configuration

• Source mode: Auto
• Repository: Unavailable
• Branch: Unavailable
• File: No configuration file from this repository applied
• Commit: Unavailable
• Layers: Organization UI
• Review profile: CHILL
• Plan: Advanced
• Run ID: 4086292a-89c3-4da2-922e-24059cd1d2e3

Commits

Reviewing files that changed from the base of the PR and between a664f5c and db9bedf.

Files selected for processing: 15

• CHANGELOG.md
• KNOWN_ISSUES.md
• docs/src/data_loader/data_loader.md
• docs/src/reduced_order_modeling/losses.md
• src/GeometricMachineLearning.jl
• src/data_loader/parametric_data_loader.jl
• src/data_loader/train.jl
• src/loss/losses.jl
• src/parametric.jl
• src/pullbacks/zygote_pullback.jl
• test/data_loader/parametric_data_loader.jl
• test/helpers/parametric_affine.jl
• test/loss/parametric_loss.jl
• test/parametric.jl
• test/runtests.jl

Included review availability: This review used your included allowance. Your plan provides up to 1 included review per hour; 0 remain after this review.

[codecov]

Codecov Report

❌ Patch coverage is 96.15385% with 3 lines in your changes missing coverage. Please review.
✅ Project coverage is 92.32%. Comparing base (a664f5c) to head (c5409c5).

Files with missing lines	Patch %	Lines
src/parametric.jl	92.30%	3 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #357      +/-   ##
==========================================
+ Coverage   92.06%   92.32%   +0.26%     
==========================================
  Files          51       53       +2     
  Lines        1109     1186      +77     
==========================================
+ Hits         1021     1095      +74     
- Misses         88       91       +3     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[michakraus]

JuliaGNI/GeometricMachineLearning.jl PR #357 — Add the parametric core: apply_parametric, ParametricLoss, ParametricDataLoader

Verdict: request changes (posted as a comment review)
Reviewed: 1f6670cf5f55a6d1966f4da60ed1aadb4af352aa
CI: the three red jobs are not caused by this PR. Each fails the same way on main at a664f5c9 (run 40200448244, the same day).

• Downgrade - ubuntu-latest: the job fails at the resolver, before any test runs. forcedeps check failed for GPUArraysCore 0.2.1 (bound 0.2.0), Zygote 0.7.5 (0.7.0), ChainRulesCore 1.25.1 (1.0.0), NNlib 0.9.30 (0.8.0) and KernelAbstractions 0.9.38 (0.9.0). main shows the same five lines. The PR does not change Project.toml. NeuralNetworkParameters resolves to its floor, 0.5.0 (forcedeps check passed). So this job gives no test signal at the floors. I ran a local floor check instead (see below).
• Julia pre (1.14.0-alpha1) and Julia nightly (1.15.0-DEV.132): GPUCompiler does not precompile (UndefVarError: libLLVMExtra not defined in LLVM.API), so the Metal test dependency is skipped and the run stops before any test. main has the same error.
• All other jobs pass, including Julia min and Julia 1 on all three OSes, Doctests, Documentation and Metal.

Findings

#	severity	file:line	claim	evidence
1	correctness-risk	KNOWN_ISSUES.md:135-149 (K29); PR body, Unresolved and Pre-PR verification	K29 says that the batch gradient's Union{Nothing, NetworkParameters} also shows on main, and that main's ZygotePullback(FeedForwardLoss()) on Chain(Dense(2, 2, identity)) with a 2 × 1 × 8 tensor gives Any[Union{Nothing, NetworkParameters}]. That does not reproduce. The base tree and the head both infer that gradient concretely. The Union comes from this branch: the per-sample map over eachcol in src/parametric.jl:86-101. The location field (src/pullbacks/zygote_pullback.jl:37-43) points at the (input, output) method, which this PR does not change. The decision to defer the instability as "not caused by the branch" therefore rests on a wrong premise. Please correct K29 (cause, location, and remove "the same inference shows on origin/main"). Then decide again whether to fix the instability or accept it.	Base tree a664f5c9 and head, Julia 1.13.2, Base.return_types(back, (typeof(v),)) for ZygotePullback(FeedForwardLoss()): 2x1x8, 2x8 and vector inputs, NeuralNetwork parameters and hand-built parameters, Float32 and Float64. Every case gives Any[NetworkParameters{Float64 or Float32, (:L1,), …}]. Head, ParametricLoss on a batch of 8: Any[Union{Nothing, NetworkParameters}]; one sample: concrete. Variants of the batch method (map + stack, a comprehension over zip, a map over one iterable with a fixed μ, a plain W * x mapped over eachcol) all give the Union. So the cause is a Zygote map over the columns with a closure that captures ps. It is not specific to the apply_parametric closure.
2	quality	src/parametric.jl:8,9,11	Three of the four _flatten_system_parameters methods have no caller. Only (::Type{T}, ::NamedTuple) is reached, from concatenate_array_with_parameters. The NullParameters methods cannot be reached, because concatenate_array_with_parameters(x, ::NullParameters) returns early. The remaining method only forwards to flatten(T, μ). Add these methods in the PRs that call them.	grep -rn _flatten_system_parameters src test docs ext finds only the definitions and the calls at :20 and :24. Codecov reports 3 lines of the patch with no coverage.
3	quality	src/parametric.jl:23-33	No code in src/ calls concatenate_array_with_parameters. Only the test helper calls it, and only the vector method (test/helpers/parametric_affine.jl:29). Both matrix methods are unused. The body itself says the matrix–vector method has an O(k²) pullback, and its plan-wide decision forbids that pattern. Shipping it now leaves that pattern for the next layer to call. The simpler file keeps two methods: concatenate_array_with_parameters(x::AbstractArray, ::NullParameters) = x and concatenate_array_with_parameters(x::AbstractVector, μ::NamedTuple) = vcat(x, first(flatten(eltype(x), μ))). Drop _flatten_system_parameters (finding 2) and the two matrix methods.	grep as in finding 2. The O(k²) figures (5931/16625/59554 B per sample) are from the PR body. I did not measure them again.
4	quality	src/data_loader/parametric_data_loader.jl:45-50	The constructor rejects a μ whose element type differs from the data's, even though the forward pass already converts it to the element type of the state. An integer system parameter (n = j,) and a Float32 μ with Float64 data are both an ArgumentError. concatenate_array_with_parameters([1.0, 2.0], (n = 3,)) returns [1.0, 2.0, 3.0]. test/loss/parametric_loss.jl trains Float32 data with Float64 μ values and keeps Float32. The check protects no path the suite exercises. Convert in the constructor or remove the check.	Probe at head: ParametricDataLoader(rand(2,5,3), [(n = j,) for j in 1:3]) gives ArgumentError: the data have the element type Float64 and a system parameter has the element type Int64. The same happens for Float32 μ values. concatenate_array_with_parameters(Float32[1, 2], (a = 0.5,)) gives Float32[1.0, 2.0, 0.5].
5	quality	src/pullbacks/zygote_pullback.jl:45-58	The two new methods have identical bodies. One method whose third slot is <:Union{OptionalParameters, AbstractVector} does the same work. That matches the signature of ParametricLoss (src/loss/losses.jl:252), which already takes the union. The comment at :28-31 then names three shapes, not four methods.	Read the diff.
6	nit	PR body	The body is written for the agents that built the branch, not for a reader on GitHub. Examples: "Decided (advisor, 2026-10-25, high)", "Decided (arbitrator, 2026-10-25, high), trigger stop after PASS", "the finish builder opens the PR", "the known-issues rule permits that without a waiver", and the heading "Decisions taken without the user". This review treats these as claims, not as settled decisions. Findings 1 and 3 hold whatever was decided.	Read the PR body.
7	nit	test/data_loader/parametric_data_loader.jl:113,121,124	occursin("3", message), occursin("2", message) && occursin("3", message) and occursin("0", message) match one digit anywhere in the message. Match the phrase instead, as :116 does with "prediction window of 2".	Read the file.

Verified and fine

• Between 24be74bd (where the body reports the suite runs) and the head, only KNOWN_ISSUES.md changed (diff stat of the two commits).
• The batch pullback allocates linearly in the batch size. I measured in a fresh process with --check-bounds=auto, Julia 1.13.2, NeuralNetworkParameters 0.5.4, AbstractNeuralNetworks 0.9.3 and Zygote 0.7.14, as the minimum of 5 readings after two warm-up calls. Bytes per sample at k = 64/256/1024: array batch 2674/2618/2636, (q, p) batch 4889/4610/4549. These agree with the body's figures. The machine had a load average of 3 to 4.
• Compat floor: NeuralNetworkParameters 0.5.0 exports flatten, parameter_eltype and parameterlayout (read in its src/NeuralNetworkParameters.jl:46,55,59). I used a scratch environment with NeuralNetworkParameters 0.5.0, AbstractNeuralNetworks 0.9.2 and GeometricOptimizers 0.9.0 pinned, on Julia 1.13.2 with --check-bounds=auto. The three new test files passed 3013 of 3013 there.
• The shape check on μ is real. ParameterLayout's == compares the children as NamedTuples, so it compares keys and their order too (NeuralNetworkParameters src/layout.jl:92-94).
• ParametricDataLoader(::EnsembleSolution) gets one μ per trajectory even when the problem uses one shared parameter set. EnsembleProblem stores parameters with the axes of ics (GeometricEquations src/problems/ensemble_problem.jl:111,142-162).
• K28 reproduces: typeof(AbstractNeuralNetworks._norm((q = rand(Float32,2,1,3), p = rand(Float32,2,1,3)))) is Float64 on 0.9.3.
• Test.detect_ambiguities(GeometricMachineLearning; recursive = false) finds 0. Aqua.test_piracies passes. ExplicitImports reports no improper explicit import of the three new names. JuliaFormatter (sciml) passes on the 11 changed .jl files.
• closingIssuesReferences is empty. The body writes "relates to #207".
• Existing review thread: CodeRabbit flagged the reuse of K22 and K23. It is resolved and fixed in 1f6670c. K28 and K29 each appear once, and the CHANGELOG.md uses neither.

Not checked

• The time cost of the Union in finding 1. The body reports equal gradient times for the closure and a callable struct, but that compares two forms that are both unstable.
• The full Pkg.test(), at the head or at the floor (not run, by instruction). The floor run uses Julia 1.13.2, not 1.12. It covered the three new test files only. Zygote, ChainRulesCore, NNlib, KernelAbstractions and GPUArraysCore were not forced to their floors.
• Device arrays. apply_parametric puts a flattened host μ below each column of the batch, so a GPU batch is probably not supported. The PR makes no claim about GPUs.
• The O(k²) figures for the matrix–vector concatenate_array_with_parameters.

[michakraus]

Resolution of the review findings

Commit c51e94e resolves the findings of the review posted above, as follows. One finding needs a decision by the author.

#	severity	action	resolution
1	correctness-risk	partly fixed; needs a decision	Reproduced on Julia 1.13.2. On a664f5c9 and on this branch, the ZygotePullback(FeedForwardLoss()) gradient on Chain(Dense(2, 2, identity)) with a 2 × 1 × 8 batch infers as a concrete NetworkParameters{T, (:L1,), …} in Float32 and Float64. Only the ParametricLoss batch path infers as Union{Nothing, NetworkParameters}; one sample is concrete. c51e94e corrects K29: its location is now src/parametric.jl:87-101, and it no longer says that main shows the Union. The Union itself is not fixed. The branch causes it, and the known-issues convention fixes such a defect in its own PR instead of recording it. The deferral in the PR body rests on the claim about main, which does not hold. Options: (a) fix it in this PR, for example with a ChainRulesCore rule for the batch method, or with another formulation that Zygote infers concretely (the review found four variants that do not); (b) accept it as a known issue, with K29 as corrected.
2	quality	not fixed	The port keeps _flatten_system_parameters as one unit. The parametric layers, ForcingLayer and the generalized HNN call the other methods in the parts that follow this one. If those parts do not land, the unreachable methods go.
3	quality	not fixed; partly disputed	The matrix methods have no caller in src/, but test/parametric.jl:18-43 tests them directly, so the test helper is not the only caller. The part that this PR implements specifies all three forms of the concatenation, with tests. The O(k²) pullback of the matrix–vector method is not on the path of apply_parametric, which maps over columns.
4	quality	not fixed	A deliberate rule of this port: the data and every μ of a loader have one element type, which is #207's check, kept. The docstring states it. Relaxing it changes the constructor's contract, which is the author's decision.
5	quality	not fixed	The comment above the methods in src/pullbacks/zygote_pullback.jl states the file's convention: one method per shape, not a union over the shapes. The two methods follow it.
6	nit	not fixed	The PR body is the author's text. This comment does not change it. The statements in its "Unresolved" section about K29 are corrected by finding 1 above.
7	nit	fixed in c51e94e	The three assertions now match the phrase that carries the number: "sequence length of 3", "hold 3 trajectories" with "there are 2 system parameters", and "hold 0 trajectories". test/data_loader/parametric_data_loader.jl passes, 2930 of 2930, on Julia 1.13.2.

Not merged. CI runs for c51e94e are in progress.

[coderabbitai]

Caution

Some comments are outside the diff and can’t be posted inline due to GitHub limitations.

⚠️ Outside diff range comments (2)

Minor · Reject empty batches in check_parametric_batch. · parametric.jl:87-92

src/parametric.jl:87-92
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win

Reject empty batches in check_parametric_batch.

A d × 1 × 0 input with an empty parameter vector passes the current size check. The mapping produces an empty outputs collection, and reduce(hcat, outputs) then throws because no initial value is provided. The (q, p) batch method has the same failure. The public API documents that a batch must contain at least one sample.

Suggested fix

-    input_size[2] == 1 && input_size[3] == length(μs) || throw(DimensionMismatch(
+    input_size[2] == 1 && input_size[3] > 0 && input_size[3] == length(μs) || throw(DimensionMismatch(

Minor · Reject unsupported non-Chain models before selecting the default parametric… · train.jl:34-37

src/data_loader/train.jl:34-37
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win

Reject unsupported non-Chain models before selecting the default parametric loss.

The train! documentation says that any network can use a ParametricDataLoader, but ParametricLoss and apply_parametric dispatch only on Chain. AbstractNeuralNetworks permits valid non-Chain models, so such a network can reach the default loss and fail with a MethodError. Constrain and document the default parametric path as Chain-only while preserving explicit custom losses.

Suggested fix

- and any network on a [`ParametricDataLoader`](@ref) on [`ParametricLoss`](@ref)`()
+ and any `Chain`-backed network on a [`ParametricDataLoader`](@ref) on [`ParametricLoss`](@ref)`()

-default_loss(::NeuralNetwork, ::ParametricDataLoader, ::Batch) = ParametricLoss()
+function default_loss(nn::NeuralNetwork, ::ParametricDataLoader, ::Batch)
+    nn.model isa Chain || throw(ArgumentError(
+        "`train!` with a `ParametricDataLoader` requires a Chain-backed network; " *
+        "pass an explicit compatible loss for another model."))
+    ParametricLoss()
+end

---

Review info

Run configuration

• Source mode: Auto
• Repository: Unavailable
• Branch: Unavailable
• File: No configuration file from this repository applied
• Commit: Unavailable
• Layers: Organization UI
• Review profile: CHILL
• Plan: Advanced
• Run ID: e189e7c6-6d12-4583-b935-f7f308282388

Commits

Reviewing files that changed from the base of the PR and between db9bedf and c51e94e.

Files selected for processing: 2

• KNOWN_ISSUES.md
• test/data_loader/parametric_data_loader.jl

Files skipped from review as they are similar to previous changes: 1

• KNOWN_ISSUES.md

Included review availability: This review used your included allowance. Your plan provides up to 1 included review per hour; 0 remain after this review.

[michakraus]

Finding 1, fixed in c5409c5

The author chose to fix the Union in this PR.

Cause. Zygote's adjoint of map adds the gradients of the mapped function from init = nothing (Zygote/src/lib/array.jl, ∇map). The batch methods of apply_parametric mapped a closure that holds ps, so the gradient of ps inferred as Union{Nothing, NetworkParameters}.

Fix. The batch methods map their samples through map_samples(f, ps, samples...). Its ChainRulesCore.rrule pulls each sample back with Zygote.pullback. It converts the cotangent of ps with NeuralNetworkParameters' public map_cotangent(storage_gradient, ps, …) and adds the results as NetworkParameters, from the first sample on. A (q, p) sample is stacked into one vector inside the map, so that its cotangent reaches Zygote as an array. The rule uses only public names, so the ExplicitImports gate needs no new exception. The closure f holds only the Chain, whose trainable parameters are in ps, so the rule gives f no gradient.

Variants that did not work: a broadcast with Ref(ps) (still the Union); map of apply_parametric with fill(ps, k) (the ProjectTo of NetworkParameters refuses ChainRules' Tangent); an rrule through rrule_via_ad (gives Union{Nothing, …}, because the ChainRules bridge allows a zero tangent).

Evidence, on Julia 1.13.2. The fix and the code at c51e94e were compared in one process.

• Array batch: the gradient infers as NetworkParameters{T, (:L1,), …} in Float32 and Float64; it was Union{Nothing, NetworkParameters}.
• The output, the W and b gradients and the input gradient are bitwise equal to those of c51e94e, for the array and the (q, p) form.
• The new testset "the gradient of ParametricLoss infers concretely" in test/parametric.jl asserts @inferred on the pullback for one sample and for an array batch, in Float32 and Float64. It passes on the fix and throws on the batch at c51e94e. test/parametric.jl passes on Julia 1.12.7 too, with 70 of 70 tests.
• test/parametric.jl, test/loss/parametric_loss.jl, test/data_loader/parametric_data_loader.jl and test/quality/{aqua,reachability,exports}.jl pass, with 3043 of 3043 tests. test/quality/jet.jl passes with 52 tests and its 1 known broken test.

What stays. The (q, p) form still infers as the Union when the loss is ParametricLoss. The cause is AbstractNeuralNetworks' _compute_loss on a (q, p) named tuple, and main at a664f5c shows the same for a plain function of ps. With a plain sum of squares, the (q, p) batch infers concretely. This is K30, under Upstream. K29 is no longer in KNOWN_ISSUES.md.

Not merged.