Repository navigation
Add bilinear and low-rank velocity-parity Stiefel layers and trace loss - #368
Conversation
|
Warning Review limit reachedYou've used all free OSS reviews for now. Wait for the free limit to reset to keep reviewing this public repository. Next included review available in 25 minutes. View limit details
📝 Summary
Merge Risk:
 |
| Check name | Status | Explanation |
|---|---|---|
| Description check |  |
The description directly explains the new bilinear and low-rank Stiefel layers, trace loss, HOSVD initialization, documentation, and validation covered by the changeset. |
| Title check | |
The title clearly and concisely identifies the main additions: bilinear and low-rank velocity-parity Stiefel layers and the trace loss. |
| Docstring Coverage | |
No eligible functions found in the changed files. Skipping docstring coverage check. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documentat… |
| 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. |
Full details: Docstring Coverage
-
Explanation
No eligible functions found in the changed files. Skipping docstring coverage check. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documentation-only edits and deletions within surviving functions are included; unchanged diff context and fully deleted functions are excluded. Analyzed 0 functions across 0 files. (4 skipped: 4 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.
Comment @coderabbitai help to get the list of available commands.
b144023 to
8e28bb2
Compare
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
-
-
- 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:
98d7d193-0237-4c4a-8862-a9b8249be2d9
-
CHANGELOG.mdKNOWN_ISSUES.mdProject.tomldocs/inventories/GeometricOptimizers.tomldocs/make.jldocs/src/architectures/transformer.mddocs/src/data_loader/data_loader.mddocs/src/data_loader/snapshot_matrix.mddocs/src/introduction.mddocs/src/layers/stiefel_layer.mddocs/src/layers/symmetric_particle_layers.mddocs/src/tutorials/symplectic_autoencoder.mddocs/src/tutorials/sympnet_tutorial.mdscripts/Project.tomlscripts/reproduction/symplectic_autoencoders/online_sympnet.jlscripts/utilities/parametric_data_helpers.jlsrc/GeometricMachineLearning.jlsrc/architectures/forced_sympnet.jlsrc/architectures/parametric_resnet.jlsrc/architectures/symmetric_symplectic_autoencoder.jlsrc/architectures/sympnet.jlsrc/architectures/transformer_integrator.jlsrc/architectures/transformer_neural_network.jlsrc/data_loader/batch.jlsrc/data_loader/data_loader.jlsrc/data_loader/parametric_data_loader.jlsrc/data_loader/tensor_assign.jlsrc/data_loader/train.jlsrc/layers/bilinear_stiefel_layer.jlsrc/layers/canonical_slot_layer.jlsrc/layers/canonical_sort_layer.jlsrc/layers/low_rank_stiefel_layer.jlsrc/layers/pairwise_symplectic_layer.jlsrc/loss/dlra_trace_loss.jlsrc/loss/losses.jlsrc/parametric.jltest/Project.tomltest/architectures/parametric_resnet.jltest/data_loader/accuracy.jltest/data_loader/autoencoder_keyword.jltest/data_loader/batch.jltest/data_loader/batch_functor.jltest/data_loader/batch_index_set.jltest/data_loader/data_loader.jltest/data_loader/draw_batch_for_tensor_test.jltest/data_loader/parametric_data_loader.jltest/data_loader/train.jltest/integration/docstrings/data_loader.jltest/integration/qp_dispatch.jltest/layers/bilinear_and_low_rank_stiefel.jltest/layers/symmetric_particle_layers.jltest/optimizers/optimizer.jltest/parametric.jltest/quality/aqua.jltest/quality/exports.jltest/quality/jet.jltest/runtests.jl
-
- test/integration/docstrings/data_loader.jl
- test/data_loader/draw_batch_for_tensor_test.jl
- test/data_loader/accuracy.jl
- test/data_loader/autoencoder_keyword.jl
- docs/src/architectures/transformer.md
- src/data_loader/tensor_assign.jl
- test/data_loader/batch_index_set.jl
- test/data_loader/batch_functor.jl
- src/data_loader/parametric_data_loader.jl
- src/data_loader/data_loader.jl
- src/data_loader/batch.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.
-
| # Spatial mode-1 unfolding: Nx × (Nv * n_snaps) | ||
| M_x = hcat(reshape(Fp_w, Nx, half_Nv * n_snaps), reshape(Fm_w, Nx, half_Nv * n_snaps)) | ||
| svd_x = svd(M_x; full = true) | ||
| _copy_stiefel!(ps.Φx, @view(svd_x.U[:, 1:Mx])) | ||
|
|
||
| # Parity-velocity mode-2 unfolding: (Nv ÷ 2) × (2 * Nx * n_snaps) | ||
| M_v = hcat( | ||
| reshape(permutedims(Fp_w, (2, 1, 3)), half_Nv, Nx * n_snaps), | ||
| reshape(permutedims(Fm_w, (2, 1, 3)), half_Nv, Nx * n_snaps)) | ||
| svd_v = svd(M_v; full = true) | ||
| _copy_stiefel!(ps.Φv, @view(svd_v.U[:, 1:Mv])) |
There was a problem hiding this comment.
Do not use svd(...; full = true) on the snapshot unfoldings. It allocates a square matrix the size of the column count.
svd(M_x; full = true) returns the full right factor V. Its size is (Nv * n_snaps) × (Nv * n_snaps). For M_v, the full right factor has size (2 * Nx * n_snaps) × (2 * Nx * n_snaps). Only U[:, 1:Mx] and U[:, 1:Mv] are used.
Example with Nx = Nv = 64 and n_snaps = 200:
M_vhas 25600 columns.- Its full
Vneeds about 5.2 GB. - The snapshot data itself is only about 6.5 MB.
With realistic snapshot counts, solve! therefore runs out of memory or stalls.
full = true only matters in the rank-deficient case, where mx or mv exceeds the number of columns. That case is covered at test lines 657-670. Use a thin SVD for the normal case. In the rank-deficient case, complete the basis with QR. With this change, memory stays proportional to size(M).
🐛 Proposed fix
-
--- "a/src/layers/bilinear_stiefel_layer.jl" +++ "b/src/layers/bilinear_stiefel_layer.jl" @@ -450,17 +450,22 @@ F_plus, F_minus = _parity_split(F_snapshots) Fp_w, Fm_w = _apply_metric_forward_3d(F_plus, F_minus, l.metric_factors) +function _leading_left_singular_vectors(M::AbstractMatrix, m::Int) + F = svd(M) # thin SVD: U is N × min(N, ncols) + size(F.U, 2) >= m && return F.U[:, 1:m] + # Rank-deficient unfolding: complete span(U) to an orthonormal N × m basis. + return qr(F.U).Q * Matrix{eltype(F.U)}(I, size(M, 1), m) +end + # Spatial mode-1 unfolding: Nx × (Nv * n_snaps) M_x = hcat(reshape(Fp_w, Nx, half_Nv * n_snaps), reshape(Fm_w, Nx, half_Nv * n_snaps)) - svd_x = svd(M_x; full = true) - _copy_stiefel!(ps.Φx, @view(svd_x.U[:, 1:Mx])) + _copy_stiefel!(ps.Φx, _leading_left_singular_vectors(M_x, Mx)) # Parity-velocity mode-2 unfolding: (Nv ÷ 2) × (2 * Nx * n_snaps) M_v = hcat( reshape(permutedims(Fp_w, (2, 1, 3)), half_Nv, Nx * n_snaps), reshape(permutedims(Fm_w, (2, 1, 3)), half_Nv, Nx * n_snaps)) - svd_v = svd(M_v; full = true) - _copy_stiefel!(ps.Φv, @view(svd_v.U[:, 1:Mv])) + _copy_stiefel!(ps.Φv, _leading_left_singular_vectors(M_v, Mv)) return ps end
Define the helper at file scope, outside
solve!.
| inverse_layer(::CanonicalSortLayer{M, Axis}) where {M, Axis} = InvertCanonicalSortLayer{M}() | ||
| inverse_layer(::InvertCanonicalSortLayer{M}) where {M} = CanonicalSortLayer{M, 1}() |
There was a problem hiding this comment.
inverse_layer does not round-trip for Axis != 1.
InvertCanonicalSortLayer{M} does not store Axis. Line 185 therefore maps it back to CanonicalSortLayer{M, 1}. inverse_layer(inverse_layer(CanonicalSortLayer{4, 2}())) returns CanonicalSortLayer{4, 1}(). The result sorts along the wrong coordinate row, and no error is raised.
SymmetricAutoencoderLayer avoids this by building the encoder's sort layer from arch.sort_axis. Other callers that apply inverse_layer to a decoder chain do not get this protection. The pairwise layers round-trip, and the tests only cover the default axis (lines 59-61 in the test file).
Document this limitation in the InvertCanonicalSortLayer docstring. Alternatively, carry Axis as a second type parameter. If you document it, add one sentence to the docstring:
📝 Proposed docstring note
-
-When called on `(q, p)` or a flattened array without a stored `perm`, it acts as the identity on the canonical Weyl chamber. +When called on `(q, p)` or a flattened array without a stored `perm`, it acts as the identity on the canonical Weyl chamber. +Note that `inverse_layer(::InvertCanonicalSortLayer)` returns `CanonicalSortLayer{M, 1}`, because the sort axis is not stored. Construct `CanonicalSortLayer{M, Axis}` explicitly for `Axis != 1`.
| const N_INDEPENDENCE_TOLERANCE = 4096 | ||
| n_independent(a::Real, b::Real) = abs(a - b) <= N_INDEPENDENCE_TOLERANCE |
There was a problem hiding this comment.
Make n_independent a one-sided check. The absolute per-sample tolerance fails because of fixed overhead, not because of superlinear growth.
The pipeline failure is at Line 854. CI reports these per-sample values:
alloc_full_bwd2 / B2≈ 76168alloc_full_bwd1 / B1≈ 80315
The gap is about 4147 bytes, which exceeds N_INDEPENDENCE_TOLERANCE = 4096. The large batch allocates less per sample than the small batch, so allocation does not grow superlinearly.
The cause is fixed overhead. Total allocation follows c + s * B. The per-sample value is then c / B + s. Between B1 = 16 and B2 = 64, the per-sample gap is c * (1/16 - 1/64). That gap exceeds 4096 bytes when c is above about 87 KB. The Zygote pullback for the full DLRA chain reaches that size.
The same test is fragile at Lines 812, 818, 832, 840, and 846. Any growth in fixed overhead, such as a new Julia or Zygote version, fails these assertions too.
The test should only reject per-sample growth with batch size. Every call site passes the B2 value first, so a one-sided comparison fixes all assertions.
💚 Proposed fix
-
--- "a/test/layers/bilinear_and_low_rank_stiefel.jl" +++ "b/test/layers/bilinear_and_low_rank_stiefel.jl" @@ -10,8 +10,10 @@ include("../helpers/eltypes.jl") const N_INDEPENDENCE_TOLERANCE = 4096 -n_independent(a::Real, b::Real) = abs(a - b) <= N_INDEPENDENCE_TOLERANCE +# `a` is the per-sample allocation at the larger batch and `b` at the smaller batch. +# Fixed overhead amortizes, so only reject per-sample growth with batch size. +n_independent(a::Real, b::Real) = a <= b + N_INDEPENDENCE_TOLERANCE function _random_stiefel(rng, ::Type{T}, N::Int, m::Int) where {T <: Real} A = randn(rng, T, N, m)
Also applies to: 848-854
Source: Pipeline failures
0bf0e73 to
1e47d17
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
-
-
- 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:
bd024aeb-21a5-42ab-82a5-844dce9d0517
-
CHANGELOG.mddocs/make.jlsrc/GeometricMachineLearning.jltest/runtests.jl
-
- docs/make.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.
-
|
|
||
| - **`CanonicalSortLayer`, `InvertCanonicalSortLayer`, `PairwiseSymplecticLayer`, `PairwiseSymplecticLayerQ`, `PairwiseSymplecticLayerP`, `InversePairwiseSymplecticLayerQ`, `InversePairwiseSymplecticLayerP`, `CanonicalSlotLayer`, `InverseCanonicalSlotLayer`, and `SymmetricSymplecticAutoencoder` add particle symplectic layers that are invariant or equivariant under relabelling of the particles, and a weight-tied symplectic autoencoder with three particle frontends.** All ten names are exported. `CanonicalSortLayer(M, d, axis = 1)` makes `CanonicalSortLayer{N, M, Axis}` with `N = 2Md`, sorts `M` particles in ``\mathbb{R}^d`` along coordinate row `axis` of `q` using a stable merge sort (`MergeSort`), applies the same permutation ``\sigma \in S_M`` to both `q` and `p`, and returns `(q = q_sorted, p = p_sorted, perm = σ)` on `(q, p)` `NamedTuple` inputs of shape `(d, M)` or `(d, M, B)` (or the sorted state on flattened `2Md` vectors and `(2Md, B)` matrices). `InvertCanonicalSortLayer(M, d)` makes `InvertCanonicalSortLayer{N, M}` and applies ``\sigma^{-1}`` (`invperm`) to restore the original particle ordering; `inverse_layer` pairs it with `CanonicalSortLayer`. `CanonicalSortLayer`, `InvertCanonicalSortLayer`, `CanonicalSlotLayer` and `InverseCanonicalSlotLayer` carry the phase-space dimension of the flattened state as their first type parameter, `N = 2Md`, so `input_dimension` and `output_dimension` return the length of that state, and their inputs are checked against the fixed `d`, with the first dimension of an input exactly its input dimension. `PairwiseSymplecticLayerQ(M, d, upscale, σ)` and `PairwiseSymplecticLayerP(M, d, upscale, σ)` apply ``S_M``-equivariant canonical shears combining a self term ``K_{\mathrm{self}}^T \sigma(K_{\mathrm{self}} u_i + b_{\mathrm{self}})`` and a skew-symmetric pairwise relative potential term ``\frac{1}{M}\sum_{j=1}^M K_{\mathrm{pair}}^T \tau(K_{\mathrm{pair}}(u_i - u_j))`` with ``\tau(x) = \sigma(x) - \sigma(-x)``, conserving total canonical momentum (or position) when ``K_{\mathrm{self}} = 0`` and pairing with `InversePairwiseSymplecticLayerQ` and `InversePairwiseSymplecticLayerP` through `inverse_layer`. `CanonicalSlotLayer(M, d, NSlots)` (`M >= 2`, `1 <= NSlots <= M`, `CanonicalSlotLayer{N, NOut, M, NSlots}` with `NOut = 2 NSlots d`) first sorts the particles along coordinate row 1 of `q` with `CanonicalSortLayer`, so its slots are invariant under a relabelling of particles with distinct first coordinates. It then separates the center-of-mass mode ``(Q_{\mathrm{cm}}, P_{\mathrm{cm}}) = \frac{1}{\sqrt{M}}\sum_{i=1}^M (q_i, p_i)`` from the ``M - 1`` zero-sum Helmert relative modes, sorts the ``M - 1`` relative modes canonically by descending squared phase-space norm ``\|q_{\mathrm{rel}, k}\|_2^2 + \|p_{\mathrm{rel}, k}\|_2^2`` (`MergeSort`), and projects the sorted relative modes onto ``N_{\mathrm{slots}} - 1`` slots via a learned `StiefelManifold` parameter ``\Phi_{\mathrm{rel}} \in \mathrm{St}(N_{\mathrm{slots}} - 1, M - 1)`` (`NSlots == 1` keeps only the center-of-mass slot with no parameter leaf). `InverseCanonicalSlotLayer(M, d, NSlots)` takes the same type parameters, with input dimension `NOut` and output dimension `N`, and lifts the slots back symplectically via ``\Phi_{\mathrm{rel}}`` and the transpose Helmert matrix. `SymmetricSymplecticAutoencoder(M, d, reduced_dim; frontend = :sort | :pairwise | :slot, ...) <: SymplecticCompression` combines one of the three particle frontends with the weight-tied `WeightTiedSymplecticAutoencoder` core. With `:sort` and `:slot` the latent state is invariant under a relabelling of particles with distinct coordinates along the sorting axis (`sort_axis` for `:sort`, coordinate 1 for `:slot`). With `:pairwise` only the frontend is permutation-equivariant, because the core is not, so the latent state is not invariant. Right-inversion ``\mathrm{encoder}(nn)(\mathrm{decoder}(nn)(\xi)) = \xi`` holds to machine precision for every ``\xi`` with `:pairwise`; with `:sort` and `:slot` it holds only for ``\xi`` whose decoded state lies in the canonical chamber, where the sorting permutations on the decoded state are the identity, because the decoder does not map into that chamber in general. The tests are `test/layers/symmetric_particle_layers.jl`. | ||
|
|
||
| - **`BilinearStiefelLayer`, `TransposeBilinearStiefelLayer`, `LowRankStiefelLayer`, `DecodeLowRankStiefelLayer`, `FastDLRATraceLoss`, and `solve!` on `BilinearStiefelLayer` add two-sided velocity-parity Stiefel frontends, a gauge-invariant dynamical low-rank Stiefel frontend and decoder, a fast low-rank trace loss, and Tucker-2 HOSVD initialization.** The five new types are exported. `BilinearStiefelLayer(Nx, Nv, mx, mv; metric_factors = nothing)` (`Nx >= 1`, even `Nv >= 2`, `1 <= mx <= Nx`, `1 <= mv <= Nv ÷ 2`) allocates `(Φx = StiefelManifold(...), Φv = StiefelManifold(...))` with ``\Phi_x \in \mathrm{St}(m_x, N_x)`` and ``\Phi_v \in \mathrm{St}(m_v, N_v / 2)``, splits an input array ``F`` of shape `(Nx, Nv)` or `(Nx, Nv, B)` into velocity-even ``F^+ = \frac{1}{\sqrt{2}}(F_{:, 1:N_v/2} + F_{:, N_v:-1:N_v/2+1})`` and velocity-odd ``F^- = \frac{1}{\sqrt{2}}(F_{:, 1:N_v/2} - F_{:, N_v:-1:N_v/2+1})`` components (or accepts a pre-split `(q, p)` `NamedTuple` of shape `(Nx, Nv ÷ 2)` or `(Nx, Nv ÷ 2, B)`), applies optional Cholesky mass factors `(Rx, Rv)` when `metric_factors` is supplied, and projects ``Q_{\mathrm{red}} = \Phi_x^T F^+ \Phi_v`` and ``P_{\mathrm{red}} = \Phi_x^T F^- \Phi_v`` to return vectorized `(q = vec(Q_red), p = vec(P_red))` of half-dimension ``n = m_x m_v``; `TransposeBilinearStiefelLayer` (`inverse_layer(l)`) is its exact right-inverse. `LowRankStiefelLayer(Nx, Nv, mx, mv; r_out = mx * mv, metric_factors = nothing)` takes a DLRA `NamedTuple` `(X = X, S = S, V = V)` of arbitrary rank ``r_{\mathrm{in}} \ge 1`` and evaluates ``Q_{\mathrm{red}} = (\Phi_x^T X) S ((V^+)^T \Phi_v)`` and ``P_{\mathrm{red}} = (\Phi_x^T X) S ((V^-)^T \Phi_v)`` in ``O(r_{\mathrm{in}}(m_x N_x + m_v N_v))`` operations without forming ``X S V^T``, matching `BilinearStiefelLayer` on ``F = X S V^T`` identically and remaining invariant under ``O(r_{\mathrm{in}}) \times O(r_{\mathrm{in}})`` gauge transformations, while `DecodeLowRankStiefelLayer` (`inverse_layer(l)`) computes the truncated SVD of ``[Q_{\mathrm{red}}\;\; P_{\mathrm{red}}]`` of rank ``r_{\mathrm{eff}} = \min(r_{\mathrm{out}}, m_x, 2 m_v)`` and QR re-orthonormalizes the lifted factors to return `(X = X_hat, S = S_hat, V = V_hat)` on the Stiefel manifold. `FastDLRATraceLoss()` evaluates ``\|X S V^T - \hat{X} \hat{S} \hat{V}^T\|_F^2`` in ``O(r^2(N_x + N_v))`` operations via Gram-matrix traces clamped at `zero(T)`. `solve!(nn, F_snapshots)` and `solve!(l, ps, F_snapshots)` initialize `Φx` and `Φv` on a `BilinearStiefelLayer` from the top-``m_x`` spatial and top-``m_v`` parity-velocity left singular vectors of `F_snapshots`. The tests are `test/layers/bilinear_and_low_rank_stiefel.jl`. |
There was a problem hiding this comment.
Correct the rank-complexity claims.
LowRankStiefelLayer performs rank-space products such as Hx * S, adding terms such as O(m_x r_in^2). FastDLRATraceLoss forms products such as S' * Gx11 * S, adding O(r^3) work, or O(B r^3) for batches. Update the bounds or state a fixed-rank assumption. (raw.githubusercontent.com)
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #368 +/- ##
==========================================
+ Coverage 95.32% 95.45% +0.12%
==========================================
Files 58 61 +3
Lines 1820 2264 +444
==========================================
+ Hits 1735 2161 +426
- Misses 85 103 +18 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
github-project-automation
Bot
moved this to Done
in Geometric Machine Learning Development
Summary
Add two-sided velocity-parity bilinear Stiefel layers, dynamical low-rank approximation (DLRA) Stiefel layers, a fast low-rank trace loss, and Tucker-2 HOSVD initialization to
GeometricMachineLearning:BilinearStiefelLayer(Nx, Nv, mx, mv; metric_factors = nothing)andTransposeBilinearStiefelLayer(inverse_layer(l)): split a 2D(Nx, Nv)or 3D(Nx, Nv, B)phase-space array into velocity-even and velocity-odd components(q, p)NamedTuple), apply optional Cholesky mass factors(Rx, Rv), and project onto(q = vec(Q_red), p = vec(P_red))of half-dimensionLowRankStiefelLayer(Nx, Nv, mx, mv; r_out = mx * mv, metric_factors = nothing)andDecodeLowRankStiefelLayer(inverse_layer(l)): encode a factored DLRA tuple(X = X, S = S, V = V)of arbitrary rankBilinearStiefelLayeron_thin_qr_cgs) to return(X = X_hat, S = S_hat, V = V_hat)on the Stiefel manifold.FastDLRATraceLoss(): evaluatezero(T).solve!(l::BilinearStiefelLayer, ps, F_snapshots)andsolve!(nn::NeuralNetwork, F_snapshots): initializeps.Φxandps.Φvvia truncated Tucker-2 HOSVD (svd(...; full = true)) on spatial and parity-velocity snapshot unfoldings.docs/src/layers/stiefel_layer.mdandtest/layers/bilinear_and_low_rank_stiefel.jl.Stacked on #367 (Part D).
Pre-PR verification
Verdict: clear to open — 0 unresolved
Diff: 9 files, +1862/−1, 3 commits base b144023...head 0bf0e73
Scope: full seven points
Fixed in this pass
Unresolved — the PR does not open until these are answered
Pre-existing, adjacent — fix in this change?
Checked and clean
src/layers/bilinear_stiefel_layer.jl:1-497,src/layers/low_rank_stiefel_layer.jl:1-354,src/loss/dlra_trace_loss.jl:1-125,src/GeometricMachineLearning.jl:251-275,docs/make.jl:199-203,docs/src/layers/stiefel_layer.md:1-23,test/layers/bilinear_and_low_rank_stiefel.jl:1-857,test/runtests.jl:28,CHANGELOG.md:151) and checked againsttest/layers/layer_structure.jl,test/quality/exports.jl,test/quality/reachability.jl, anddocs/check_references.jl. All type constructors use::Type{T}, all 6 new docstrings appear indocs/src/layers/stiefel_layer.mdwith no blank/comment line before the documented definition, 3D batched contractions use batched BLAS GEMM (reshape/permutedimsandtensor_tensor_mul/tensor_transpose_tensor_mul),DecodeLowRankStiefelLayeruses functional classical Gram-Schmidt (_thin_qr_cgs) for Zygote reverse-mode AD compatibility, andsolve!usessvd(...; full = true)to handle rank-deficient unfoldings.AbstractNeuralNetworks,NeuralNetworkParameters,GeometricOptimizers, andReducedComplexityModeling; velocity-parity bilinear Stiefel projections, DLRA factor layers, andFastDLRATraceLossdo not duplicate dependency APIs.src/layers/bilinear_stiefel_layer.jl,src/layers/low_rank_stiefel_layer.jl, andsrc/loss/dlra_trace_loss.jl. Every method added to an imported function (AbstractNeuralNetworks.inverse_layer,AbstractNeuralNetworks.parameterlength,AbstractNeuralNetworks.initialparameters,ReducedComplexityModeling.solve!) dispatches on a package-owned type (BilinearStiefelLayer,TransposeBilinearStiefelLayer,LowRankStiefelLayer,DecodeLowRankStiefelLayer, orNeuralNetwork{..., <:Chain{<:Tuple{BilinearStiefelLayer, Vararg}}}).promote_type(eltype(x), eltype(ps.Wx), ...)/inv(sqrt(T(2)))/zero(T)arithmetic across all forward, decode, and loss paths, plus@inferredassertions intest/layers/bilinear_and_low_rank_stiefel.jlon all constructors,inverse_layer,initialparameters, and 2D/3D calls.fatou lint --no-configon all 7 changed.jlfiles (src/layers/bilinear_stiefel_layer.jl,src/layers/low_rank_stiefel_layer.jl,src/loss/dlra_trace_loss.jl,src/GeometricMachineLearning.jl,test/layers/bilinear_and_low_rank_stiefel.jl,test/runtests.jl,docs/make.jl) with 0 findings on new or modified hunks. Confirmed@viewusage on array slices and two-batch-size scaling assertions (B1 = 16,B2 = 64) intest/layers/bilinear_and_low_rank_stiefel.jl:804-855.src/layers/bilinear_stiefel_layer.jl,src/layers/low_rank_stiefel_layer.jl,src/loss/dlra_trace_loss.jl,docs/src/layers/stiefel_layer.md, andtest/layers/bilinear_and_low_rank_stiefel.jl; all prose is in the present tense and matches the implementation.CHANGELOG.md: grepped all changed files for temporal/history phrases (previously,used to,was renamed,no longer,as of,before this,originally,recently,has been changed to,going forward,from now on) with 0 hits outsideCHANGELOG.md, and verified the## [Unreleased]entry atCHANGELOG.md:151.Not checked
Aqua.test_piracies): Kaimon live session creation was disabled because the Kaimon server had stalled sessions (test/quality/aqua.jlrunsAqua.test_all(GeometricMachineLearning; ambiguities = false)in CI).code_typed): Kaimon live session creation was disabled (@inferredcovers all new layers and loss paths intest/layers/bilinear_and_low_rank_stiefel.jl).@allocated): Julia cannot run directly in the terminal sandbox and Kaimon session creation was disabled.Changelog: entry present
Release: minor bump — adds exported
BilinearStiefelLayer,TransposeBilinearStiefelLayer,LowRankStiefelLayer,DecodeLowRankStiefelLayer, andFastDLRATraceLoss(0.x repository, covered by the pending 0.9.0 release inCHANGELOG.md)Follow-up measurement: none
Critic
Verdict: PASS on
0bf0e734(Round: verify 3)BilinearStiefelLayerandTransposeBilinearStiefelLayer(inverse_layer),initialparameters, velocity-parity split/recombine, optionalmetric_factors, right-inverse identity totest/layers/bilinear_and_low_rank_stiefel.jl(BilinearStiefelLayer constructor, guards, parametersandBilinearStiefelLayer forward, inverse_layer, parity isometry, Kronecker symplecticity)LowRankStiefelLayerandDecodeLowRankStiefelLayer(inverse_layer), arbitrary input rankBilinearStiefelLayertotest/layers/bilinear_and_low_rank_stiefel.jl(LowRankStiefelLayer and DecodeLowRankStiefelLayer)FastDLRATraceLoss()zero(T), matching dense Frobenius norm totest/layers/bilinear_and_low_rank_stiefel.jl(FastDLRATraceLoss)solve!(nn::NeuralNetwork, F_snapshots)andsolve!(l::BilinearStiefelLayer, ps, F_snapshots)Tucker-2 HOSVD initialization matching optimal subspace projection on rank-$(m_x, m_v)$ synthetic data totest/layers/bilinear_and_low_rank_stiefel.jl(solve! Tucker-2 HOSVD initialization on BilinearStiefelLayer)docs/src/,@inferred, Aqua, ExplicitImports, GML test suite,CHANGELOG.md, and batched Zygote differentiation / allocation scalingdocs/src/layers/stiefel_layer.md,CHANGELOG.md,test/layers/bilinear_and_low_rank_stiefel.jl(Batched Zygote differentiation and allocation scaling)Known issues added
none
Decisions taken without the user
cayley_direct!,DifferentiatedCayleyTransport) solves itslu!/ldiv!. On Metal it runs underallowscalar(false): the reviewer of the pull request measured both against the host and got differences of about 1e-7 and 8e-7 inFloat32. On aJLArrayit raises "Scalar indexing is disallowed.", because JLArrays has nolu. No clause of this plan runs it on a device, and part A has no device test (it would go intest/devices/metal.jl). A part that puts a direct Stiefel method on a device measures it on Metal through Kaimon, and never reads device support from aJLArraytest.r"^…$"), not a substring. A phrase such as "a matrix of size" also occurs in Base's ownDimensionMismatchof a matrix product, so a substring test passes when the guard is removed and a later operation raises.DecodeLowRankStiefelLayermaps_decode_dlra_2dovereachcol(vcat(reshape(Q_red, Mx * Mv, B), reshape(P_red, Mx * Mv, B)))rather than indexingQ_red[:, :, b]andP_red[:, :, b]insidemap(1:B), avoidinggetindexpullbacks and_dlra_trace_loss_3dusestensor_transpose_tensor_mul,tensor_tensor_mul, and whole-array trace reductionssum(M .* permutedims(G, (2, 1, 3)); dims = (1, 2))withsum(max.(zero(T), t1 .- T(2) .* tc .+ t2)), deleting_batched_at_b,_batched_a_b,_batched_a_bt, and the duplicateChainRulesCore.rrule(::typeof(_dlra_trace_loss_3d), ...).