You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
Repository navigation
Add MultiModalSymplecticAutoencoder, multimodal losses, and HRedSys reference shift and VJP - #370
Implements Part F of the multi-modal symplectic autoencoders plan in GeometricMachineLearning:
MultiModalSymplecticAutoencoder, encode, decode, and translate:
MultiModalSymplecticAutoencoder(frontends::NamedTuple, core::SymplecticCompression) <: SymplecticCompression combines a non-empty NamedTuple of modality frontends (individual layers or Chains of symplectic particle layers) whose canonical output dimension 2n_m matches core.full_dim with a shared SymplecticCompression core (SymplecticAutoencoder, WeightTiedSymplecticAutoencoder, or SymmetricSymplecticAutoencoder).
encode(nn, modality, x) (with modality a Symbol or Val(modality)) applies frontends[modality] followed by encoder(core), decode(nn, modality, ξ) applies decoder(core) followed by inverse_layer(frontends[modality]), and translate(nn, from, to, x) computes decode(nn, to, encode(nn, from, x)).
ChainRulesCore.rrule methods for encode and decode (each for a Symbol and a Val modality) differentiate the active frontend and the core with rrule_via_ad and give zero tangents for the other frontends, so Zygote.gradient with respect to network parameters works across all modalities, Chain frontends included.
Multi-modal training losses (CharbonnierRelativeLoss, IsometricRolloutLoss, CrossModalCycleLoss):
CharbonnierRelativeLoss(; ε = 1e-3, λ_lat = 1e-2) evaluates the exact-zero Charbonnier relative reconstruction loss plus latent consistency penalty, dispatching low-rank (X, S, V) inputs through FastDLRATraceLoss without forming dense Nx × Nv arrays.
HRedSys reference-state shifting, :vjp pullback mode, and preallocated buffers:
Extends HRedSys, build_v_reduced, and build_f_reduced with optional reference states q_ref and p_ref (converted to Vector{eltype(ics.q)}) and pullback_mode::Symbol = :jacobian | :vjp.
:vjp evaluates $(\nabla_\xi \Psi^{\mathrm{dec}}(\xi))^T \mathbb{J}_{2N}^T [v(z); f(z)]$ in a single reverse-mode pullback Zygote.pullback(decoder, ξ) and falls back to ForwardDiff.jacobian! when state vectors carry ForwardDiff.Dual elements during implicit integrator Newton iterations.
Preallocates work buffers (ξ_buf, v_intermediate, f_intermediate, z_bar_buf, jac_result) and a per-element-type IdDict{DataType, Any} cache (ReducedVFCache) for ForwardDiff.Dual evaluations.
Pre-PR verification
Fixed in this pass
#
point
file:line
what was wrong
what changed
—
—
—
none
—
Unresolved
#
point
file:line
claim
evidence
—
—
—
none
—
Pre-existing, adjacent
#
severity
file:line
claim
evidence
Checked and clean
1 · correctness, quality, concision: run-tests.jl passed 645 of 645 across the six files (test/reduced_system/reference_shift_and_vjp.jl, test/reduced_system/reduced_system.jl, test/architectures/multimodal_symplectic_autoencoder.jl, test/quality/aqua.jl, test/quality/exports.jl, test/quality/reachability.jl), from one run of run-tests.jl on Julia 1.13.3. JuliaFormatter (style = "sciml") and fatou lint are clean on the changed source files, and test-layout.jl --check reports no violation. Two mutants are caught by test/reduced_system/reference_shift_and_vjp.jl: evaluate_vf_and_compute_∇Ψ returning its cache buffers instead of copies, and a floating-point state of another element type than the cache's taking the per-type cache instead of raising an ArgumentError. A one-sided _add_ref! that shifts the wrong component is caught by the same file, and the original _check_eltype (which promoted the components of a mixed state) is caught by test/architectures/multimodal_symplectic_autoencoder.jl.
2 · functionality duplicated from dependencies: The code reuses FastDLRATraceLoss, tensor_transpose_tensor_mul, tensor_tensor_mul, inverse_layer, DiffResults.JacobianResult, ForwardDiff.jacobian!, and Zygote.pullback.
3 · type piracy: The Aqua checks and the ExplicitImports checks of test/quality/aqua.jl pass (17/17).
4 · type instability: The test files assert @inferred for encode, decode, translate, and the three losses in Float32 and Float64. Base.return_types of evaluate_vf_and_compute_∇Ψ is concrete for Float64 and for ForwardDiff.Dual input.
5 · avoidable allocations: One warm call of v_reduced with 2n = 4 allocates, as the minimum of 5 calls, in two fresh processes with --check-bounds=auto on Julia 1.13.3: 4448 / 8160 / 29728 bytes with :jacobian and 27968 / 29376 / 37248 bytes with :vjp, at 2N = 12 / 40 / 200. The allocations grow with N. evaluate_vf_and_compute_∇Ψ returns copies of its cache buffers, so that no caller holds memory of the cache. The batched _frob_sq_norm allocates 316.01 MiB before and 3.01 MiB after for Float32 with Nx = Nv = 10000, rank 32 and 128 slices (minimum of 7 warm calls, fresh process, --check-bounds=auto, Julia 1.13.3).
6 · stale comments, present tense: No added line in src/, test/, or docs/src/ matches previously, originally, no longer, used to, now , or formerly.
7 · history only in CHANGELOG.md: The same search finds no historical phrase; the CHANGELOG.md entry is under ## [Unreleased].
Not checked
No Documenter build and no run on Julia 1.12 in the local checks (CI covers both); no full Pkg.test(); no mutant for the batched _frob_sq_norm copy change, whose test checks value, gradient and non-mutation.
Release: minor bump — adds MultiModalSymplecticAutoencoder, encode, decode, translate, CharbonnierRelativeLoss, IsometricRolloutLoss, CrossModalCycleLoss, and HRedSys reference-state shifting and :vjp pullback support prior to 1.0.
Critic
The builder-critic loop passed a436a12 (Round: verify 3 on 7f21517..a436a12, base 9ef9680). The review of this pull request and the fix commits after it changed the code. The table gives the state of each earlier defect at this head.
Earlier defects
#
defect
status at this head
evidence
B1 (R1)
Zygote.gradient with respect to network parameters failed on MultiModalSymplecticAutoencoder when any frontend was a Chain, and CharbonnierRelativeLoss, IsometricRolloutLoss, and CrossModalCycleLoss accepted only NeuralNetwork and failed under Zygote.gradient(ps -> loss(nn.architecture, ps, ...), nn.params)
fixed
encode and decode on a MultiModalSymplecticLayer have ChainRulesCore.rrule methods that differentiate the active frontend and the core with rrule_via_ad and give zero tangents for the other frontends. The three losses have methods for (model, ps::NamedTuple, ...). The testset "Analytical Zygote gradients match central finite differences in Float64" and the Chain-frontend testset of test/architectures/multimodal_symplectic_autoencoder.jl pass.
B2 (R1)
HRedSys did not convert q_ref and p_ref to the state element type eltype(ics.q)
fixed
The HRedSys constructor converts with Vector{T_state}(q_ref) and Vector{T_state}(p_ref). test/reduced_system/reference_shift_and_vjp.jl checks that Float64 references give Float32 fields for Float32 states.
B3 (R1)
HRedSys allocation test in test/reduced_system/reference_shift_and_vjp.jl ran inside @safetestset rather than in a cold julia --startup-file=no --check-bounds=yes subprocess
not at this head
Neither new test file has an allocation test. The figures in "Checked and clean" come from a separate cold script.
B1 (verify 1)
MultiModalSymplecticAutoencoder used ChainRulesCore.@non_differentiable _core_layer_keys(Any) and @non_differentiable _frontend_inverse_layer(Any), violating the spec's ban on custom ChainRules rules
not at this head
src/architectures/multimodal_symplectic_autoencoder.jl has four ChainRulesCore.rrule methods (encode and decode, each for a Symbol and a Val modality) and one @non_differentiable (_check_modality). src/loss/multimodal_losses.jl has two @non_differentiable (_check_traj_len and _check_eltype).
Known issues added
K40, K41
Decisions taken without the user
Decided (arbitrator, 2026-10-30, high): Replace the ad-hoc _apply_frontend(fe::Chain, ...) pre-flattening and runtime keys(ps_core)[1:M] slicing with compositional helper functions that preserve Chain's nested parameter NamedTuple structure and split MultiModalSymplecticAutoencoder encoder/decoder parameter NamedTuple keys via Val(M) and Base.structdiff, so that Zygote reverse-mode AD produces well-typed non-nothing parameter cotangents for Chain frontends and all three loss types on both NamedTuple and NeuralNetwork wrappers without custom ChainRulesCore.@non_differentiable rules.
Decided (arbitrator, 2026-10-26): The direct (I - A)\(I + A) Cayley retraction and its analytical differential and pullback are implemented for AbstractMatrix and AbstractVecOrMat via standard lu(I - A) without scalar indexing. On CPU arrays, this runs natively without fallback. On GPU arrays whose backend does not support batch-less lu(I - A) or where a caller prefers the series, the retraction dispatches or falls back to the truncated Neumann series. Because the Vlasov–Poisson training uses small reduced Stiefel dimensions ($n \le 32$, $m_x, m_v \le 16$), Stiefel retraction steps run on CPU or small device matrices where lu is fast.
Decided (arbitrator, 2026-10-26; user, 2026-10-26): Each part adds its own new test file(s) and includes them in test/runtests.jl, and runs only the new and directly touched test files during the builder-critic loop.
Included review availability: This review used your included allowance. Your plan provides up to 1 included review per hour; 0 remain after this review.
📝 Summary
Summary by CodeRabbit
New Features
Added multimodal symplectic autoencoders with encode, decode, and translate operations across modalities.
Added reconstruction, rollout, and cross-modal cycle losses for training and evaluation.
Reduced-system modeling now supports reference-state shifts in encoding, decoding, and reduced evaluations, with selectable Jacobian or VJP pullbacks.
Documentation
Expanded library documentation to cover multimodal autoencoder operations and losses, and updated the changelog to reflect these additions.
The pull request adds a multimodal symplectic autoencoder with modality-specific operations and three losses. It also extends HRedSys with reference-state shifts, Jacobian and VJP pullback modes, and reusable typed caches.
Changes: Multimodal symplectic compression
Layer / File(s)
Summary
Architecture contract and construction src/architectures/multimodal_symplectic_autoencoder.jl, src/GeometricMachineLearning.jl, docs/src/architectures/symplectic_autoencoder.md, test/architectures/multimodal_symplectic_autoencoder.jl
Adds the multimodal architecture, validates frontends against the shared core, initializes frontend and core paths, and exports and documents the API.
Modality operations and gradients src/architectures/multimodal_symplectic_autoencoder.jl, test/architectures/multimodal_symplectic_autoencoder.jl
Adds modality-specific encode, decode, and translate operations. Tests cover frontend and core combinations, round trips, batched inputs, and gradient behavior.
Multimodal losses and validation src/loss/multimodal_losses.jl, docs/src/reduced_order_modeling/losses.md, test/architectures/multimodal_symplectic_autoencoder.jl, test/runtests.jl
Adds CharbonnierRelativeLoss, IsometricRolloutLoss, and CrossModalCycleLoss, with tests for loss values, invalid inputs, and gradients.
Reduced-system reference shifts and pullbacks
Layer / File(s)
Summary
Reference state and cache setup src/reduced_system/reduced_system.jl, test/reduced_system/reference_shift_and_vjp.jl
Adds reference and pullback options to HRedSys constructors, validates reference dimensions and pullback modes, and stores reusable work buffers and typed caches.
Updates reduced-field builders to apply reference offsets and select Jacobian or VJP evaluation. Tests compare pullback results and check cache behavior.
Reference-aware encoding and integration src/reduced_system/reduced_system.jl, test/reduced_system/reference_shift_and_vjp.jl
Adds reference-aware encode and decode methods and uses them in integration and error calculations. Tests cover shifted equilibria, Hamiltonians, and element-type conversion.
Change records and test registration CHANGELOG.md, KNOWN_ISSUES.md, test/runtests.jl
Records the new multimodal and reduced-system APIs. The known-issue entry describes HRedSys cache-related fields and pullback_mode that source code does not read.
sequenceDiagram
participant HRedSys
participant ReducedVFCache
participant Decoder
participant FullVectorFields
participant Zygote
HRedSys->>ReducedVFCache: retrieve typed work buffers
HRedSys->>Decoder: decode reduced state
Decoder->>FullVectorFields: evaluate at reference-shifted state
FullVectorFields->>Zygote: provide vector-field cotangent
Zygote->>HRedSys: return reduced-coordinate pullback
Loading
Possibly related PRs
JuliaGNI/GeometricMachineLearning.jl#363: Adds weight-tied symplectic autoencoders and analytic inverse layers used by the multimodal architecture’s core paths.
Large batched loss evaluations may use avoidable memory, and the cache documentation needs correction. These are bounded follow-ups rather than merge blockers.
Pre-merge checks: 4
Check name
Status
Explanation
Title check
The title clearly and concisely summarizes the three main changes: the multimodal autoencoder, multimodal losses, and HRedSys reference shifting and VJP support.
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 and comprehensively describes the multimodal autoencoder, multimodal losses, and HRedSys changes listed in the changeset.
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.
Included review availability: This review used your included allowance. Your plan provides up to 1 included review per hour; 0 remain after this review.
The reason will be displayed to describe this comment to others. Learn more.
Minor · Functional Correctness · Quick win
Remove the res::T assertion in IsometricRolloutLoss. It throws a TypeError when the integrator or H changes the element type.
T comes only from xs[1] and ps.core. The accumulators err_lat, err_rec and err_H also take the element type of loss.integrator(z_curr) and of loss.H(z).
Take a Float32 network with a user integrator z -> A * z, where A is a Float64 matrix. Then z_curr is Float64 and res is Float64. The assertion res::T expects Float32 and throws. The same happens when a custom H returns Float64.
CharbonnierRelativeLoss and CrossModalCycleLoss do not have this problem. Their T already includes the element types of the encoded latents.
Return res without the assertion. If you want a concrete accumulator type, derive it from z1, for example promote_type(T, eltype(z1)). The integrator and H can still widen the type.
The reason will be displayed to describe this comment to others. Learn more.
Minor · Functional Correctness · Quick win
Supported by static analysis
Script executed:
#!/bin/bash# Check the layout returned by data_tensors_from_geometric_solution and existing test coverage.
rg -nP -A25 'function\s+data_tensors_from_geometric_solution' src
rg -nP -C3 '\b(reduction_error|projection_error)\s*\('test
Sources consulted:
- https://raw.githubusercontent.com/JuliaRegistries/General/master/R/ReducedComplexityModeling/Versions.toml
- https://raw.githubusercontent.com/JuliaRegistries/General/master/R/ReducedComplexityModeling/Package.toml
- GitHub - SciML/ModelOrderReduction.jl: High-level model-order reduction to automate the acceleration of large-scale simulations · GitHub — https://github.com/SciML/ModelOrderReduction.jl
- SciMLDocs/docs/src/highlevels/equation_solvers.md at main · SciML/SciMLDocs · GitHub — https://github.com/SciML/SciMLDocs/blob/main/docs/src/highlevels/equation_solvers.md
- GitHub - MORFEproject/MORFE.jl: Reduced Order Modelling via the Parametrisation Method · GitHub — https://github.com/MORFEproject/MORFE.jl
- …and 41 more
Script executed:
#!/bin/bashset -o pipefail
tree='bed0ac78af861e0b5ed0d9bf77a871aa4e2d2916'
api="https://api.github.com/repos/JuliaRCM/ReducedComplexityModeling.jl/git/trees/${tree}?recursive=1"
json=$(curl -fsSL "$api")printf'%s\n''--- matching paths ---'printf'%s'"$json"| python3 -c 'import json, sysdata=json.load(sys.stdin)for item in data.get("tree", []): if item.get("type") == "blob": print(item["path"])'| grep -Ei 'loader|solution|data|src'| head -100
printf'%s\n''--- matching source lines ---'printf'%s'"$json"| python3 -c 'import json, sysdata=json.load(sys.stdin)for item in data.get("tree", []): if item.get("type") != "blob" or not item["path"].endswith((".jl", ".md")): continue url="https://raw.githubusercontent.com/JuliaRCM/ReducedComplexityModeling.jl/bed0ac78af861e0b5ed0d9bf77a871aa4e2d2916/"+item["path"] import urllib.request try: text=urllib.request.urlopen(url).read().decode() except Exception: continue if "data_tensors_from_geometric_solution" in text: print("PATH", item["path"]) for i,line in enumerate(text.splitlines(), 1): if "data_tensors_from_geometric_solution" in line or (i > 1 and "function data_tensors" in text.splitlines()[i-2]): lo=max(1,i-8); hi=min(len(text.splitlines()),i+35) for n in range(lo,hi+1): print(f"{n}: {text.splitlines()[n-1]}")'
#!/bin/bashset -o pipefail
tree='bed0ac78af861e0b5ed0d9bf77a871aa4e2d2916'
api="https://api.github.com/repos/JuliaRCM/ReducedComplexityModeling.jl/git/trees/${tree}?recursive=1"
python3 - "$api"<<'PY'import base64, json, subprocess, sys, urllib.requesttree = json.load(urllib.request.urlopen(sys.argv[1]))base = "https://api.github.com/repos/JuliaRCM/ReducedComplexityModeling.jl/git/blobs/"for item in tree.get("tree", []): path = item.get("path", "") if item.get("type") != "blob" or not path.endswith(".jl"): continue if not any(x in path.lower() for x in ("trainingdata", "data_loader", "reducedcomplexitymodeling")): continue try: blob = json.load(urllib.request.urlopen(base + item["sha"])) text = base64.b64decode(blob["content"]).decode() except Exception: continue if any(term in text for term in ("data_tensors_from_geometric_solution", "geometric_solution", "GeometricSolution", "StateVariable")): print(f"--- {path} ---") lines = text.splitlines() for i, line in enumerate(lines, 1): if any(term in line for term in ("data_tensors_from_geometric_solution", "geometric_solution", "GeometricSolution", "StateVariable")): for n in range(max(1, i - 10), min(len(lines), i + 45) + 1): print(f"{n}: {lines[n-1]}") print()PY
Repository: JuliaGNI/GeometricMachineLearning.jl
Length of output: 41545
Broadcast references over the leading data axis.
data_tensors_from_geometric_solution returns (system dimension, time steps, 1). A vector reference broadcasts against the trailing axis in Julia, so rs.q_ref and rs.p_ref create an unintended third dimension instead of shifting each system component. This makes both error metrics incorrect for non-zero references.
Reshape each reference vector to (N, 1, 1) for tensor inputs, and add a non-zero-reference test for both reduction_error and projection_error.
The reason will be displayed to describe this comment to others. Learn more.
JuliaGNI/GeometricMachineLearning.jl PR #370 — Add MultiModalSymplecticAutoencoder, multimodal losses, and HRedSys reference shift and VJP
Verdict: request changes (posted as a comment review, because the PR author is the account that posts it)
Reviewed: a436a12fbf4bcf3f5d6fdb09754d6e555a486f3e
CI: red on every release job because of this diff. Julia 1 (macOS, ubuntu, windows) and Julia min (macOS, windows) each fail with one error in test/quality/aqua.jl:101:
every name reached in GeometricOptimizers is public: Error During Test at .../test/quality/aqua.jl:101
Expression: check_all_qualified_accesses_are_public(GeometricMachineLearning; ignore = NON_PUBLIC_QUALIFIED) === nothing
NonPublicQualifiedAccessException
- `MutableDiffResult` is not public in `DiffResults` but it was imported from `DiffResults` at `.../src/reduced_system/reduced_system.jl:6:29`
Some tests did not pass: 16 passed, 0 failed, 1 errored, 0 broken.
Not from this diff: Julia pre (1.14.0-alpha2) and Julia nightly (1.15.0-DEV.258) fail to precompile GPUCompiler (UndefVarError: libLLVMExtra not defined in LLVM.API; The following 1 package failed to precompile: GPUCompiler). Downgrade fails at julia-downgrade-compat before any test runs (forcedeps check failed: Zygote resolved to 0.7.5 but lower bound is 0.7.0, and the same for ChainRulesCore, NNlib and KernelAbstractions). The same four lines fail the Downgrade job of main at 9ef9680 (run 40950534234). Both Metal jobs and Doctests pass. Julia min ubuntu, Documentation and Scripts were still running when this review was written.
Findings
#
severity
file:line
claim
evidence
1
blocker
src/reduced_system/reduced_system.jl:6
The field type DiffResults.MutableDiffResult{1, Vector{T}, Tuple{Matrix{T}}} names a non-public type. The ExplicitImports check in test/quality/aqua.jl:101 errors on all five finished release jobs. Simplest fix: make the type a parameter, ReducedVFCache{T, JR} with jac_result::JR, as HRedSys already does with JR. Do not add :MutableDiffResult to NON_PUBLIC_QUALIFIED.
The CI lines above. Locally: Base.ispublic(DiffResults, :MutableDiffResult) gives false.
2
bug
src/loss/multimodal_losses.jl:9-15
CharbonnierRelativeLoss fails on a batched DLRA input (X, S, V) with 3-D factors. encode and decode accept this input (the test file checks it at lines 564-567). _frob_sq_norm calls transpose on the 3-D X and V.
CharbonnierRelativeLoss(; ε = 1e-2, λ_lat = 5e-2)(nn, :dlra, (X = X_b, S = S_b, V = V_b)) with size(X_b) == (6, 2, 3) gives transpose not defined for Array{Float64, 3}. Consider using permutedims for higher-dimensional arrays.
initializer is a named keyword, so kwargs... never holds it. The frontends get the caller's initializer and the core always gets its own default, without a warning. CodeRabbit's inline comment is right.
With the parameter-free frontend CanonicalSortLayer(2, 2): NeuralNetwork(core, CPU(), Float64; initializer = ZeroInitializer()) throws OrthonormalizationFailure, but NeuralNetwork(MultiModalSymplecticAutoencoder((particles = fe,), core), CPU(), Float64; initializer = ZeroInitializer()) builds, and its 24 core leaves are not all zero.
4
quality
PR body, sections "Critic" and "Checked and clean"
The body describes code that is not at this head. (a) It says git grep "non_differentiable|ChainRules|rrule|frule" over the three source files "returns 0 matches" and that the spec bans custom ChainRules rules. The head defines four ChainRulesCore.rrule methods (multimodal_symplectic_autoencoder.jl:403,439,494,532) and two @non_differentiable (multimodal_symplectic_autoencoder.jl:25, multimodal_losses.jl:185). (b) It names _split_core_params(ps, ::Val{M}), Base.structdiff and _apply_frontend_chain. None of them exists in src/. (c) It cites a cold-subprocess allocation test at test/reduced_system/reference_shift_and_vjp.jl:194-216. Those lines are the end of the ForwardDiff.Dual cache test and the start of an equilibrium test, and the file has no @allocated. (d) It says ExplicitImports is clean, but finding 1 shows that it fails. Either the code breaks the spec that the body cites, or the body is stale. The body must state which one.
The body claims "verified size-independent allocations across 2N = 12 and 2N = 40". The claim does not reproduce. The allocations of a warm v_reduced call grow with N, because the decoder pass allocates. Against main, the buffers save about 20-35 % of the bytes.
--check-bounds=auto, warm, minimum of 5 calls, 2n = 4, one fresh process each, on a loaded machine. main default: 4736 / 8688 / 31520 B at 2N = 12 / 40 / 200. This head, :jacobian: 3728 / 6304 / 20624 B. This head, :vjp: 27968 / 29376 / 37248 B.
The docstring offers :vjp as an alternative to :jacobian but does not give its cost. At every measured size, :vjp is one to two orders of magnitude slower per v_reduced call, and v_reduced and f_reduced each run their own decoder pullback. State the trade in the docstring, or show the system size at which :vjp gets ahead.
Same script, one run on a loaded machine. :jacobian: 0.74 / 0.89 / 2.83 µs. :vjp: 95 / 101 / 104 µs at 2N = 12 / 40 / 200.
HRedSys gets five buffer fields, the type parameter JR and the field pullback_mode. Nothing in src/ reads them. The cache lives in the v_reduced and f_reduced closures. Only hasfield assertions in the test read the buffer fields. Remove the fields, JR, and the hasfield tests (test/reduced_system/reference_shift_and_vjp.jl:177-182).
git grep -e ξ_buf -e jac_result -e dual_cache -e pullback_mode a436a12 -- src test docs: no read access outside the constructor and the tests.
8
quality
src/loss/multimodal_losses.jl:192,212
IsometricRolloutLoss asserts res::T with T taken only from xs[1] and ps.core. A Float32 network with an integrator or H that returns Float64 throws TypeError. CodeRabbit's inline comment is right. Promote T to include the integrator's and H's output types, or remove the assertion.
Float32 network, step = z -> A64 * z with A64::Matrix{Float64}: TypeError: in typeassert, expected Float32, got a value of type Float64. With H = z -> 0.5 * sum(abs2, z) and an identity step, the same error.
9
quality
test/runtests.jl:66,74
Both new files are in core, but test/architectures/multimodal_symplectic_autoencoder.jl takes more than the 60 s limit of that group, even warm in CI. The PR body lists this as "pre-existing, adjacent". These files are new in this PR, so the placement comes from this diff. Move the file to slow.
The macOS Julia 1 log, in a warm process: MultiModalSymplecticAutoencoder and multimodal losses … 182 182 1m12.1s and HRedSys reference-state shifting and VJP pullback mode … 267 267 30.8s. The body itself gives 198.2 s and 133.7 s cold.
The ForwardDiff.Dual test checks only size and all(isfinite, …), so a wrong Dual-path Jacobian passes. Compare it with central differences of v_reduced, or with the :vjp result applied column by column.
v_reduced and f_reduced of one HRedSys share one mutable ReducedVFCache (ξ_buf, jac_result, v_intermediate, f_intermediate) and return views into it. On main, each call allocated its own arrays. Two tasks that integrate the same HRedSys now race. At a minimum, state in the docstring that the object is not safe to share between threads.
The default cache of build_v_reduced and build_f_reduced is ReducedVFCache(Float64, …). A Float32 caller therefore always takes the IdDict path with a type assertion, never the direct buffer. HRedSys passes eltype(ics.q). The standalone builders could take the type from decoder's parameters.
Read the lines. _get_vf_cache returns the base cache only when T === T0 (line 22).
13
nit
src/reduced_system/reduced_system.jl:373-380
HRedSys always passes vectors (zeros by default), so the nothing branch of h_reduced never runs from HRedSys. The shifted branch allocates two N-vectors on each call, even for a zero reference.
(model::MultiModalSymplecticLayer)(x, ps), and therefore nn(x), silently round-trips through the first modality in keys(frontends). No docstring states this. Throw an ArgumentError that points to encode and decode, as Chain(::MultiModalSymplecticAutoencoder) does at lines 73-79, or document the behaviour.
Test labels and comments cite "Clause 1/2/3", "Mutant M1", "catches O3/O5", "catches M2/M3" and "Proposition 13.1". These names come from a plan document that is not in the repository, so a reader cannot follow them. Name the property that each test checks.
CodeRabbit's inline comment about broadcast over the leading axis is wrong. Julia broadcasting extends trailing singleton dimensions, so qp.q .- rs.q_ref shifts each component of an (N, T, 1) tensor. However, no test calls reduction_error or projection_error with a non-zero reference. Add one.
size(rand(3, 5, 1) .- [10.0, 20.0, 30.0]) == (3, 5, 1), and the shift acts on axis 1. With q_ref, p_ref given: projection_error = 0.0935, reduction_error = 0.0935. Without: 0.853, 0.853. Neither run raised an error.
Verified and fine
Sign and block layout of the :vjp path. z̄ = [-f; v] gives v_reduced = g[n+1:2n] and f_reduced = -g[1:n], which matches J_{2n} (∇Ψ)ᵀ J_{2N}ᵀ [v; f] and the :jacobian branch (lines 276-280, 317, 323, 356, 362).
The custom rrules for encode and decode give correct parameter gradients on the paths that the test only checks for finiteness: a Chain frontend (PairwiseSymplecticLayerQ and PairwiseSymplecticLayerP), a SymplecticAutoencoder core, and translate, CharbonnierRelativeLoss and CrossModalCycleLoss. Zygote against central differences (h = 1e-6) over every parameter leaf: the maximum absolute error is 4.9e-10 to 3.3e-9, with gradient magnitudes of 1.4 to 18.
The two new test files pass in CI: 182/182 and 267/267 on macOS Julia 1.
The Downgrade, pre and nightly failures do not come from this diff (see the CI section). The PR does not change Project.toml.
test-layout.jl --check on the head reports no violations.
The PR body holds no closing keyword: closingIssuesReferences is empty.
Not checked
Julia min ubuntu, Documentation and Scripts: they were still running.
Device arrays. The _narrow_* helpers (lines 194-224) assert Vector{T} and Matrix{T}, so a device array would raise a TypeError there. The probe could not reach these helpers, because NeuralNetwork(WeightTiedSymplecticAutoencoder(…), JLBackend(), Float32) already fails on main's code path (Illegal conversion of a JLArray to a Ptr in cholesky, psd_like_layer.jl:33). Is the multi-modal autoencoder meant to support device arrays?
Whether the twelve _narrow_* methods, which exist only to assert types, can be removed without loss of inference. No inference probe was run.
- ReducedVFCache takes its Jacobian result type as a parameter instead of
naming DiffResults' non-public MutableDiffResult, which failed the
ExplicitImports check on every release job.
- CharbonnierRelativeLoss accepts batched DLRA input (3-D X, S, V factors).
- NeuralNetwork(::MultiModalSymplecticAutoencoder, ...) passes its
initializer keyword to the core network as well as to the frontends.
- Tests for the batched loss and for the initializer reaching the core.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
ReducedVFCache{T, JR} holds jac_result::JR, so the source no longer names DiffResults.MutableDiffResult. check_all_qualified_accesses_are_public with the ignore list of test/quality/aqua.jl returns nothing, and test/quality/aqua.jl passes locally (17/17). On the Dual path the cache's JR is not inferred, so DiffResults.value and DiffResults.jacobian are asserted to Vector{T} and Matrix{T}; Base.return_types of evaluate_vf_and_compute_∇Ψ is concrete for both Float64 and Dual input. The Dual-path Jacobian of v_reduced agrees with central differences to 4.3e-11.
A _frob_sq_norm method for 3-D X, S, V sums the squared norms of the slices, with the package's tensor kernels. It equals the dense sum(abs2, X_b * S_b * V_b') summed over the batch; the batched loss and its Zygote gradient are finite. New tests in test/architectures/multimodal_symplectic_autoencoder.jl.
initializer is passed to the core NeuralNetwork. The default is unchanged: DefaultInitializer === GlorotUniform in the resolved AbstractNeuralNetworks, and the default call gives the same parameters as an explicit GlorotUniform(). A new test with a counting initializer checks that the core calls it.
4
quality
not fixed
The PR body is the author's text. Its "Critic" and "Checked and clean" sections do not match this head and should be rewritten or removed by the author.
5
quality
not fixed
Same: the claim of size-independent allocations in the body does not reproduce and should be removed from the body.
6
quality
needs a decision
Whether to document :vjp as the slower mode, or to compute one pullback for both v_reduced and f_reduced, is a design choice.
7
quality
needs a decision
The five buffer fields, JR and pullback_mode on HRedSys are never read in src/, but the tests assert that they exist (hasfield). Removing them changes what the tests specify.
8
quality
needs a decision
The res::T assertion turns a Float64 integrator or H with a Float32 network into a TypeError. Whether to convert, promote, or reject with an ArgumentError is an API choice.
9
quality
needs a decision
Moving the two new test files from core to slow changes the suite's grouping.
10
quality
not fixed
A stronger Dual-path test needs a chosen reference and tolerance. A central-difference comparison (agreement 4.3e-11 in Float64) is a candidate.
11
correctness-risk
needs a decision
The shared mutable cache makes v_reduced and f_reduced non-reentrant. A per-task cache or a documented restriction is a design choice.
12
nit
not fixed
Standalone builders default the cache to Float64; changing the default is part of the same API question as 8.
13
nit
not fixed
Left for the author: it does not change results, and removing the dead branch or the two allocations is a cleanup choice.
14
nit
not fixed
Documenting nn(x) is a docs choice for the author.
15
nit
not fixed
The plan labels in the new test files are left for the author to rename.
16
nit
disputed (CodeRabbit's part)
As the review said: the length-N reference broadcasts over axis 1 of the (N, T, 1) tensor, so reduction_error and projection_error are correct. The missing test with a non-zero reference remains open.
Not addressed by this push: the Downgrade, pre and nightly failures, which also fail on main and are not caused by this PR.
Included review availability: This review used your included allowance. Your plan provides up to 1 included review per hour; 0 remain after this review.
❌ Patch coverage is 88.65546% with 54 lines in your changes missing coverage. Please review.
✅ Project coverage is 94.37%. Comparing base (9ef9680) to head (c0030b4).
… fields
- IsometricRolloutLoss and CrossModalCycleLoss raise an ArgumentError that
names both element types when the data, the integrator's output or H
has another element type than the network's parameters.
- build_v_reduced and build_f_reduced default their cache to the
decoder's parameter element type; a floating-point state of another
element type than the cache's raises an ArgumentError. Dual states keep
their own cache.
- evaluate_vf_and_compute_∇Ψ returns copies, so no caller holds memory
of the cache.
- KNOWN_ISSUES.md K40: HRedSys fields that no code in src/ reads.
- Tests for each error and for results that do not alias the cache.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Follow-up: the author decided the open findings, and 775e13f carries them.
#
decision
action
7
keep the fields
kept; recorded as K40 in KNOWN_ISSUES.md; the hasfield tests stay
8
reject mixed precision with an ArgumentError that names the types
IsometricRolloutLoss checks the data, the integrator's output and the value of H; CrossModalCycleLoss checks x_a and x_b; one test per check
12
same as 8, in the builders
the default cache of build_v_reduced and build_f_reduced has the decoder's parameter element type; a floating-point state of another element type than the cache's raises an ArgumentError; ForwardDiff.Dual states keep their own cache; tested
11
never return a view into a cache
evaluate_vf_and_compute_∇Ψ returns copies; a test mutates a result and checks the next call, and checks that a later call leaves an earlier result alone, for both helpers. A mutant that returns the cache's buffers is caught. Concurrent use of one HRedSys from two tasks still shares the cache buffers; this change does not make it reentrant.
9
move to slow only if over 60 s cold
measured alone, cold: 56.5 s and 33.7 s (testset times), so both stay in core
4, 5
rewrite the "Checked and clean" and "Critic" sections
done; the allocation figures were measured again, in two fresh processes, and grow with N
Findings 6, 10 and 13–16 are unchanged from the earlier comment.
The reason will be displayed to describe this comment to others. Learn more.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to GitHub limitations.
⚠️ Outside diff range comments (2)
Minor · Check each component’s element type. · multimodal_symplectic_autoencoder.jl:2
test/architectures/multimodal_symplectic_autoencoder.jl:2 Minor · Functional Correctness · Quick win
Check each component’s element type.
For a Float64 network, a particle state with Float64 q and Float32 p passes _check_eltype because _state_eltype promotes them to Float64. CanonicalSlotLayer then raises a MethodError instead of the documented ArgumentError. Check each component in _check_eltype and add a mixed-component test. (raw.githubusercontent.com)
Minor · Handle omitted references in standalone reduced-field builders. · reduced_system.jl:44-53
src/reduced_system/reduced_system.jl:44-53 Minor · Functional Correctness · Quick win
Handle omitted references in standalone reduced-field builders.
build_v_reduced and build_f_reduced document both references as optional. HRedSys also defaults each reference independently to a zero vector. When a standalone builder receives only q_ref or only p_ref, its evaluation path calls _add_ref! with one vector and one nothing, which raises MethodError.
Add one-sided _add_ref! methods that shift only the supplied component.
Suggested fix
_add_ref!(z::AbstractVector, ::Nothing, ::Nothing, ::Int) = z
++function _add_ref!(+ z::AbstractVector, q_ref::AbstractVector, ::Nothing, N2::Int)+ @inbounds for i in 1:N2+ z[i] += q_ref[i]+ end+ z+end++function _add_ref!(+ z::AbstractVector, ::Nothing, p_ref::AbstractVector, N2::Int)+ @inbounds for i in 1:N2+ z[N2 + i] += p_ref[i]+ end+ z+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: 0e45db27-d315-4f68-a44a-3373d6a8ee7b
Commits
Reviewing files that changed from the base of the PR and between 432509b and 775e13f.
Included review availability: This review used your included allowance. Your plan provides up to 1 included review per hour; 0 remain after this review.
…acobian, drop dead branch
Document in build_v_reduced and build_f_reduced that pullback_mode = :vjp allocates
more and runs slower than :jacobian for small problems, with measured figures.
Compare the Dual-path Jacobian of the reduced field with central finite differences
under a non-zero reference shift, reject a transposed and a sign-flipped Jacobian,
and test reduction_error and projection_error against a loop-based value with a
non-zero reference shift. Remove the unreachable nothing branch of build_h_reduced.
Document that nn(x) of MultiModalSymplecticAutoencoder uses the first modality only.
Rename the plan labels in the two new test files to descriptive names.
Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
…mponent, one-sided references
- _frob_sq_norm of a batch of DLRA states reuses a factor whose element type is already T.
- _check_eltype checks every component of a state, so a mixed-precision particle state raises the ArgumentError.
- _add_ref! shifts only the supplied component when one reference is given.
- CHANGELOG: state the default cache element type for HRedSys and for the standalone builders.
Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
michakraus
left a comment
github-project-automation
Bot
moved this to Done
in Geometric Machine Learning Development
Summary
Implements Part F of the multi-modal symplectic autoencoders plan in
GeometricMachineLearning:MultiModalSymplecticAutoencoder,encode,decode, andtranslate:MultiModalSymplecticAutoencoder(frontends::NamedTuple, core::SymplecticCompression) <: SymplecticCompressioncombines a non-emptyNamedTupleof modality frontends (individual layers orChains of symplectic particle layers) whose canonical output dimension2n_mmatchescore.full_dimwith a sharedSymplecticCompressioncore (SymplecticAutoencoder,WeightTiedSymplecticAutoencoder, orSymmetricSymplecticAutoencoder).encode(nn, modality, x)(withmodalityaSymbolorVal(modality)) appliesfrontends[modality]followed byencoder(core),decode(nn, modality, ξ)appliesdecoder(core)followed byinverse_layer(frontends[modality]), andtranslate(nn, from, to, x)computesdecode(nn, to, encode(nn, from, x)).ChainRulesCore.rrulemethods forencodeanddecode(each for aSymboland aValmodality) differentiate the active frontend and the core withrrule_via_adand give zero tangents for the other frontends, soZygote.gradientwith respect to network parameters works across all modalities,Chainfrontends included.CharbonnierRelativeLoss,IsometricRolloutLoss,CrossModalCycleLoss):CharbonnierRelativeLoss(; ε = 1e-3, λ_lat = 1e-2)evaluates the exact-zero Charbonnier relative reconstruction loss plus latent consistency penalty, dispatching low-rank(X, S, V)inputs throughFastDLRATraceLosswithout forming denseNx × Nvarrays.IsometricRolloutLoss(integrator; α_lat = 1.0, α_rec = 1.0, α_H = 0.1, H = nothing)evaluates multi-step latent rollout error, decoded trajectory error, and latent Hamiltonian conservation penalty.CrossModalCycleLoss(mod_a, mod_b; β_align = 1.0, β_cycle = 0.1)evaluates paired latent alignment plus cross-modal cycle consistency.HRedSysreference-state shifting,:vjppullback mode, and preallocated buffers:HRedSys,build_v_reduced, andbuild_f_reducedwith optional reference statesq_refandp_ref(converted toVector{eltype(ics.q)}) andpullback_mode::Symbol = :jacobian | :vjp.:vjpevaluatesZygote.pullback(decoder, ξ)and falls back toForwardDiff.jacobian!when state vectors carryForwardDiff.Dualelements during implicit integrator Newton iterations.ξ_buf,v_intermediate,f_intermediate,z_bar_buf,jac_result) and a per-element-typeIdDict{DataType, Any}cache (ReducedVFCache) forForwardDiff.Dualevaluations.Pre-PR verification
Fixed in this pass
Unresolved
Pre-existing, adjacent
Checked and clean
run-tests.jlpassed 645 of 645 across the six files (test/reduced_system/reference_shift_and_vjp.jl, test/reduced_system/reduced_system.jl, test/architectures/multimodal_symplectic_autoencoder.jl, test/quality/aqua.jl, test/quality/exports.jl, test/quality/reachability.jl), from one run ofrun-tests.jlon Julia 1.13.3.JuliaFormatter(style = "sciml") andfatou lintare clean on the changed source files, andtest-layout.jl --checkreports no violation. Two mutants are caught bytest/reduced_system/reference_shift_and_vjp.jl:evaluate_vf_and_compute_∇Ψreturning its cache buffers instead of copies, and a floating-point state of another element type than the cache's taking the per-type cache instead of raising anArgumentError. A one-sided_add_ref!that shifts the wrong component is caught by the same file, and the original_check_eltype(which promoted the components of a mixed state) is caught by test/architectures/multimodal_symplectic_autoencoder.jl.FastDLRATraceLoss,tensor_transpose_tensor_mul,tensor_tensor_mul,inverse_layer,DiffResults.JacobianResult,ForwardDiff.jacobian!, andZygote.pullback.test/quality/aqua.jlpass (17/17).@inferredforencode,decode,translate, and the three losses inFloat32andFloat64.Base.return_typesofevaluate_vf_and_compute_∇Ψis concrete forFloat64and forForwardDiff.Dualinput.v_reducedwith2n = 4allocates, as the minimum of 5 calls, in two fresh processes with--check-bounds=autoon Julia 1.13.3:4448/8160/29728bytes with:jacobianand27968/29376/37248bytes with:vjp, at2N = 12/40/200. The allocations grow withN.evaluate_vf_and_compute_∇Ψreturns copies of its cache buffers, so that no caller holds memory of the cache. The batched_frob_sq_normallocates 316.01 MiB before and 3.01 MiB after for Float32 with Nx = Nv = 10000, rank 32 and 128 slices (minimum of 7 warm calls, fresh process,--check-bounds=auto, Julia 1.13.3).src/,test/, ordocs/src/matchespreviously,originally,no longer,used to,now, orformerly.CHANGELOG.md: The same search finds no historical phrase; theCHANGELOG.mdentry is under## [Unreleased].Not checked
Pkg.test(); no mutant for the batched_frob_sq_normcopy change, whose test checks value, gradient and non-mutation.Release: minor bump — adds
MultiModalSymplecticAutoencoder,encode,decode,translate,CharbonnierRelativeLoss,IsometricRolloutLoss,CrossModalCycleLoss, andHRedSysreference-state shifting and:vjppullback support prior to1.0.Critic
The builder-critic loop passed
a436a12(Round: verify 3on7f21517..a436a12, base9ef9680). The review of this pull request and the fix commits after it changed the code. The table gives the state of each earlier defect at this head.Earlier defects
Zygote.gradientwith respect to network parameters failed onMultiModalSymplecticAutoencoderwhen any frontend was aChain, andCharbonnierRelativeLoss,IsometricRolloutLoss, andCrossModalCycleLossaccepted onlyNeuralNetworkand failed underZygote.gradient(ps -> loss(nn.architecture, ps, ...), nn.params)encodeanddecodeon aMultiModalSymplecticLayerhaveChainRulesCore.rrulemethods that differentiate the active frontend and the core withrrule_via_adand give zero tangents for the other frontends. The three losses have methods for(model, ps::NamedTuple, ...). The testset "Analytical Zygote gradients match central finite differences in Float64" and theChain-frontend testset oftest/architectures/multimodal_symplectic_autoencoder.jlpass.HRedSysdid not convertq_refandp_refto the state element typeeltype(ics.q)HRedSysconstructor converts withVector{T_state}(q_ref)andVector{T_state}(p_ref).test/reduced_system/reference_shift_and_vjp.jlchecks thatFloat64references giveFloat32fields forFloat32states.HRedSysallocation test intest/reduced_system/reference_shift_and_vjp.jlran inside@safetestsetrather than in a coldjulia --startup-file=no --check-bounds=yessubprocessMultiModalSymplecticAutoencoderusedChainRulesCore.@non_differentiable _core_layer_keys(Any)and@non_differentiable _frontend_inverse_layer(Any), violating the spec's ban on custom ChainRules rulessrc/architectures/multimodal_symplectic_autoencoder.jlhas fourChainRulesCore.rrulemethods (encodeanddecode, each for aSymboland aValmodality) and one@non_differentiable(_check_modality).src/loss/multimodal_losses.jlhas two@non_differentiable(_check_traj_lenand_check_eltype).Known issues added
K40, K41
Decisions taken without the user
_apply_frontend(fe::Chain, ...)pre-flattening and runtimekeys(ps_core)[1:M]slicing with compositional helper functions that preserveChain's nested parameterNamedTuplestructure and splitMultiModalSymplecticAutoencoderencoder/decoder parameterNamedTuplekeys viaVal(M)andBase.structdiff, so that Zygote reverse-mode AD produces well-typed non-nothingparameter cotangents forChainfrontends and all three loss types on bothNamedTupleandNeuralNetworkwrappers without customChainRulesCore.@non_differentiablerules.(I - A)\(I + A)Cayley retraction and its analytical differential and pullback are implemented forAbstractMatrixandAbstractVecOrMatvia standardlu(I - A)without scalar indexing. On CPU arrays, this runs natively without fallback. On GPU arrays whose backend does not support batch-lesslu(I - A)or where a caller prefers the series, the retraction dispatches or falls back to the truncated Neumann series. Because the Vlasov–Poisson training uses small reduced Stiefel dimensions (luis fast.test/runtests.jl, and runs only the new and directly touched test files during the builder-critic loop.