[michakraus]

Part M5 of the plan that simplifies GeometricOptimizers and GeometricMachineLearning for 0.9: the layers and the architectures.

What changes

• Q/P pairs. The q and p variants of the SympNet layers, of SymplecticAttention and of LinearSymplecticAttention share one functor each, through _shift(Val(C), z, f). LinearLayerQ/LinearLayerP lose their unused type parameter. GradientLayer.second_dim is a concrete Int.
• Bool ladders. VolumePreservingAttention, the VPFF layer, SymplecticAttention, SymplecticAutoencoder and ResNetLayer encode each choice as a Symbol type parameter, built in one expression. The public Bool keywords stay.
• VPFF. VolumePreservingFeedForwardLayer is one struct with a Triangle parameter. VolumePreservingLowerLayer and VolumePreservingUpperLayer stay exported as aliases. The bias initialiser keyword is init_bias.
• MultiHeadAttention. One initialparameters. The head keys are built once, and the forward pass reads each head by its name, so a pass allocates no key.
• Architectures. One Chain method for LASympNet. One _gradient_layers for GSympNet, LinearSymplecticTransformer and SymplecticTransformer (closes known issue C19). The DataLoader constructors delegate. GSympNet.act is GSympNet.activation. One encoder/decoder path for SymplecticCompression.
• Bugs. LinearSymplecticTransformer(dim, seq_length; activation = sin) no longer raises a MethodError. The field upsacling_activation of StandardTransformerIntegrator is upscaling_activation, and the first Dense of its chain carries that activation.
• ::Type{T} in every method under src/layers/.
• PositionalEncoding allocates its encoding on the backend of its input (closes known issues B11 and K5).

The CHANGELOG.md entries under [Unreleased] give the breaking changes and their migrations.

The MultiHeadAttention draw order is unchanged (head by head; in each head, query, key, value), and no seeded pin moved.

MultiHeadAttention (Stiefel, add_connection), ClassificationLayer (average) and AbstractNeuralNetworks' Dense are outside the Bool-ladder item, as decided below: they encode other choices and keep their Bool parameters.

Tests

• test/architectures/forward_pins.jl with test/helpers/architecture_pins.jl: the outputs of the touched architectures in Float32 and Float64 on dense, SymmetricMatrix, triangular, SkewSymMatrix and Stiefel leaves, against values recorded on the base commit 4c4a8320, at 8 ULPs.
• test/integration/forward_inference.jl: @inferred on the forward pass of every exported architecture, in both precisions.
• test/layers/multi_head_attention_allocation.jl and test/layers/multi_head_attention_key_order.jl: the head keys and the head order.
• test/layers/resnet_bias_encoding.jl, test/layers/type_arguments.jl, test/layers/volume_preserving_feedforward.jl, test/layers/positional_encoding.jl (JLArray under disallowed scalar indexing), and the architecture tests of the two bugs.

Full suite at 23a9c74f, run locally: Julia 1.13.2 and Julia 1.12.7 each pass 99 testsets, 6602 pass, 1 broken (JET, tracked), 0 fail, 0 error. The Metal group passes 40/40 on both versions (Apple M4 Max). The doctests pass on 1.13.2.

Critic

## Critic — part M5, round 3 (verify)

Verdict: PASS
Branch: simplify/part-m5-2 at 23a9c74
Judged: fbea9d28..HEAD

### Evidence
| clause | domain | my command | on origin/main | on the branch | verdict |
|:--|:--|:--|:--|:--|:--|
| Docstring adjacency | ResNetLayer | Earlier `docstring_adjacency.jl`, with branch and extracted base paths | 1/1 | 1/1 | PASS |
| Documentation exists | ResNetLayer | Cold TestEnv run of `round-3/hasdoc.jl` | Reference test resolves the binding | `Docs.hasdoc`: true; 1/1 | PASS |
| Reference check | Manual references | `run-tests.jl <worktree> quality/docs_make_checks_references.jl`; `pins.jl` on origin/main | 6/6, PINS | 6/6 on 1.13.2; 6/6 on 1.12.7 | PASS |
| M-doc | Comment between docstring and struct | `mutate.jl <worktree> --check round-3/mutants.toml`, then `--warm` | Not applicable | CAUGHT; seven unresolved references; assertions 33 and 35 fail | PASS |
| ResNet behaviour | Existing tests and decided bias encoding | Named runs of `layers/resnet_tests.jl` and `layers/resnet_bias_encoding.jl` | Not re-run; executable code unchanged by this fix | 128/128; 28/28 | PASS |
| Earlier regression guards | ClassificationTransformer and MHA head order | Named tests; O9 mutant; cold `mha_key_order.jl` | Not re-run | 4/4; 16/16; O9 CAUGHT; reproducer 4/4 with zero error | PASS |
| Quality checks | Package quality units | Named runs under `quality/` | Not re-run | Aqua 17/17; JET 52 pass, 1 known broken; reachability 2/2; exports 8/8 | PASS |
| Restricted fix | Arbitrator’s permitted change | `git diff fbea9d28..HEAD`; `git diff --check fbea9d28..HEAD` | Not applicable | Only the comment moves inside the struct; no whitespace errors | PASS |
| Earlier mutants retain their verdicts | The 22 fix-1 mutants | Read the fix diff under evidence.md’s fix-round selection rule | Not applicable | No executable code or test changes; earlier verdicts retained | PASS |

### Earlier blocking defects (verify rounds)
| round | defect | reproducer | now |
|:--|:--|:--|:--|
| 1a | Delegate test misses swapped flags | O9 in `round-3/mutants.toml` | CAUGHT; two unequal-flag assertions fail |
| 1a | Inconsistent choice encoding | Earlier `probe_one_encoding.jl` | Bias assertion passes. MHA/VPA assertion still fails, but §M5 explicitly excludes that comparison |
| 1c | MHA matches heads by position | Earlier `mha_key_order.jl` | Fixed: 4/4; zero error in both precisions and shapes |
| 2 | Comment removes ResNetLayer documentation | Earlier `docstring_adjacency.jl`; reference test | Fixed: 1/1 and 6/6 |

### Blocking defects
None.

### Other defects
1. not verified — §M5:5548 — The dispatcher’s full-suite results are not available in this review. I did not run the full suite. The dispatcher must confirm both Julia versions, including `slow`, `metal`, and `doctests`.

### Fix first
none

The critic's other defect 1 is answered by the full-suite run above.

Pre-PR verification

Fixed: none needed.
Unresolved: none.
Pre-existing, adjacent (nit):

• MultiHeadAttention reads the head count from length(ps.PQ), not d.n_heads. This is new in the diff.
• ActivationLayer, LinearLayer and SymplecticAttention constructors do not coerce M to Int, where GradientLayer now does.
• A misplaced comment above the tensor method of compute_output_of_mha.

Checked and clean:

• Correctness, read in full:
  • The LASympNet Chain merge reproduces all four old layer sequences.
  • The MultiHeadAttention draw order is unchanged.
  • Zygote gradients agree with finite differences, and structured leaves are kept.
• Piracy: Aqua Piracy.hunt returns nothing.
• Type stability: code_typed on every touched layer method, Float32 and Float64, shows concrete return types and no Any call.
• Allocations: fatou lint is clean. The MultiHeadAttention forward pass allocates exactly the same bytes as the literal-key form (cold, 8 combinations × 5 repeats). PositionalEncoding allocates one matrix plus the sum.
• GPU: PositionalEncoding agrees with the host on a real Metal GPU (matrix, 3-tensor, (q, p), gradient). This supports closing B11.
• Comments and history: no stale comment and no history in prose.
• Formatting: JuliaFormatter reports all 33 changed .jl files formatted.
• Changelog: entry present.

Not checked:

• The full suite and the docs build were not run by the verifier (the full suite ran separately, above).
• Only Julia 1.13.2 was used.
• The Metal test group was not rerun.
• The allocation test fails intermittently in a warm shared REPL (cold: clean).

Unresolved

• nit: MultiHeadAttention's forward pass takes the number of heads from its parameters, length(ps.PQ), not from d.n_heads. Parameters with more heads than the layer run without an error when add_connection = false. Left as is: a change of code after the critic's PASS is not in this pull request.

Known issues added

• K20 · LagrangianNeuralNetwork names its activation field act, where every other architecture says activation.
• K21 · ClassificationLayer's initialparameters names its initialiser init_weight!, where every other layer says init_weight.

Known issues closed: B11, K5, C19.

Advisor decisions

Decided (advisor, 2026-10-19, high): check 4's last sentence reads: "encoder/decoder of a
SymplecticAutoencoder keep NonLinearSymplecticEncoder/NonLinearSymplecticDecoder, and those of
a PSDArch (SymplecticCompression) keep UnknownSymplecticEncoder/UnknownSymplecticDecoder, as
on the parent commit 4c4a8320." The branch's code stays as it is. The test asserts all three
families over REAL_ELTYPES: the SymplecticAutoencoder pair, the PSDArch pair, and a test-local
struct PlainAE <: AutoEncoder that gives UnknownEncoder/UnknownDecoder.

Decided (advisor, 2026-10-19, medium): check 3's "equals" is abs(measured - literal) ≤ 1024 B,
with k = 1000 forward passes inside one function-barrier call and literal > 0 as the control. One
rebuilt key costs at least 240 B per pass, so the mutant exceeds the tolerance by more than 230×.

Decided (advisor, 2026-10-20, high): check 2's "σ === tanh" reads "σ === Dense(4, 6, tanh).σ",
because AbstractNeuralNetworks' Dense stores Activation(tanh) = TanhActivation() (ANN
src/activation.jl:29, src/layers/dense.jl:5-8), so no tree can satisfy the literal. The test
also builds the architecture with the default upscaling_activation and transformer_dim = 6, and
asserts Chain(arch).layers[1].σ === Dense(4, 6, identity).σ. Tests catch: Chain calling
Dense(arch.sys_dim, arch.transformer_dim) without the activation picks Dense's default tanh,
and the default-activation assertion fails. A constructor that drops the keyword stores identity,
and the tanh assertion fails. Nothing changes in AbstractNeuralNetworks.

Decided (advisor, 2026-10-20, high): the MHA item's "The RNG draw order changes; re-pin any
seeded test" and the edge "The RNG draw order" predict what the sketch in gml-simplify.md §4.4
would do. They do not require that change. The item requires one initialparameters and the head
keys built once. A build that keeps the parent's draw order (head by head; in each head, query,
key, value) meets the item. It re-pins no seeded test, and the pull request says that the draw
order is unchanged and that no seeded pin moved. Tests catch: draws reordered to "PQ for every
head, then PK, then PV" in the one initialparameters make the StandardTransformerIntegrator
pins (n_heads = 2, dense MHA leaves, test/architectures/forward_pins.jl) fail at 8 ULPs.

Decided (advisor, 2026-10-20, medium): the Bool-ladder item reaches the four ladders it
lists, and each other GML layer that encodes one of the same choices. In GML, only ResNetLayer
does, for the bias. "The same kind in every layer" means a Symbol type parameter in those layers.
ResNetLayer encodes its bias as :bias/:no_bias, as the VPFF layer does, and dispatches its
initialparameters, parameterlength and forward methods on that Symbol. Its use_bias::Bool
keyword stays. MultiHeadAttention (Stiefel, add_connection) and ClassificationLayer
(average) encode other choices, have no ladder, and keep their Bool parameters. ANN Dense
is outside M5, and the pull request names the three as out of the item. Tests catch: for
use_bias in (true, false),
typeof(ResNetLayer(3; use_bias).parameters[3]) === typeof(VolumePreservingLowerLayer(3; use_bias).parameters[3]),
and ResNetLayer(3) isa ResNetLayer{3, 3, :bias, typeof(identity)}. The mutant that passes
use_bias through as a Bool fails both. In Float32 and Float64,
parameterlength(ResNetLayer(3; use_bias = false)) == 9 and == 12 with the bias.
ResNetLayer(3, tanh; use_bias = false) gives x + tanh.(W * x), with eltype === T. The
mutant that sends :no_bias to the bias method fails both.

Arbitrator decision

Decided (arbitrator, 2026-10-20, high; M5, second FAIL): Verify 2 failed at fbea9d28 on one
fix regression. The comment that fix 1 added at src/layers/resnet.jl:17 sits between the
@doc raw"""…""" and struct ResNetLayer, so Julia attaches no docstring to ResNetLayer. This
gives the 7 unresolved references in test/quality/docs_make_checks_references.jl on 1.13.2 and
1.12.7. Fix 2 moves that comment inside the struct, or deletes it, and changes nothing else.
Verify 3 judges fbea9d28..HEAD. It is the last critic round that runs without the user.
Next round: (a) Change only src/layers/resnet.jl:17: put the comment inside the struct ResNetLayer body above activation::F1, or delete it. No other source line changes. (b) Before
it reports, the builder runs quality/docs_make_checks_references.jl (it must pass 6/6) and
Docs.hasdoc(GeometricMachineLearning, :ResNetLayer). Fix 1 ran 29 files and left out this
test. (c) Test that catches the removal: test/quality/docs_make_checks_references.jl:33,35.
Input: the docs reference check. Assertion: exactly 1 unresolved reference(s) (the fixture's
own). Mutant M-doc puts a # comment line back between the closing """ and struct ResNetLayer. It must be CAUGHT, because the count goes to 7. (d) The 22 mutants of fix 1
(R1–R7, F1–F5, N3a/b/k, O6, O9, O29, O34, O35) stay CAUGHT. (e) Verify 3 closes when (b) holds,
M-doc is CAUGHT, and the dispatcher's full suite is green on 1.13.2 and on 1.12.7, with slow,
metal and doctests run. A FAIL at verify 3 stops the loop (trigger blocked).

User decision

Decided (user, 2026-10-20): a recorded vector is compared at a stated count of ULPs of the
larger magnitude, not with ==: 8 ULPs here. The pins stay recorded from the parent commit; only
the comparison changes. The CI runners showed the SAE encoder vector differing by 1 to 2 ULPs
between x86_64 and arm64, while a wrong extraction moves it by a fraction, thousands of ULPs away.

[coderabbitai]

📝 Summary

Summary by CodeRabbit

• Breaking Changes

  • Corrected public activation and initialization option names.
  • Renamed the transformer upscaling activation field for consistency.
  • Standardized layer and architecture constructor behavior.

• Bug Fixes

  • Custom activations now propagate correctly through transformer and SympNet architectures.
  • Positional encoding now supports device-resident arrays, including GPU-compatible workflows.
  • Corrected volume-preserving feed-forward mappings and bias handling.

• Reliability

  • Improved output type preservation across supported numeric types.
  • Added reproducibility checks and broader coverage for architecture outputs, attention layers, and forward-pass inference.

This PR updates layer type parameters and shared dispatch, changes attention and positional-encoding implementations, and adjusts architecture constructors and layer assembly. It adds tests for backend behavior, parameter handling, constructor mappings, output stability, and forward-pass inference.

Changes: Layer and architecture updates

Layer / File(s)	Summary
Layer contracts and transformations src/layers/sympnets.jl, src/layers/resnet.jl, src/layers/symplectic_attention.jl, src/layers/linear_symplectic_attention.jl, src/layers/psd_like_layer.jl, src/layers/volume_preserving_attention.jl, src/layers/volume_preserving_feedforward.jl, test/layers/{resnet_bias_encoding,type_arguments,volume_preserving_feedforward}.jl, test/runtests.jl, CHANGELOG.md, KNOWN_ISSUES.md	Layer dispatch now shares Q/P transformation methods. GradientLayer stores its dimension as Int; ResNetLayer uses symbolic bias parameters; and lower and upper volume-preserving feed-forward layers use a shared concrete type. Tests cover these contracts and parameter behavior.
Attention and positional encoding src/layers/multi_head_attention.jl, src/layers/positional_encoding.jl, src/layers/transformer.jl, test/layers/{multi_head_attention_allocation,multi_head_attention_key_order,positional_encoding}.jl, test/runtests.jl, CHANGELOG.md, KNOWN_ISSUES.md	Multi-head attention initializes and retrieves projections by named head keys. Positional encoding allocates on the input backend and fills the encoding through broadcast. Tests cover head ordering, allocations, and device-array behavior.
Architecture construction and validation src/architectures/{autoencoder,linear_symplectic_transformer,standard_transformer_integrator,symplectic_autoencoder,symplectic_transformer,sympnet,transformer_neural_network}.jl, test/architectures/*, test/helpers/architecture_pins.jl, test/integration/{documented_defaults,forward_inference}.jl, test/runtests.jl, KNOWN_ISSUES.md	Architecture constructors and chain builders now use shared layer construction. GSympNet and StandardTransformerIntegrator use the corrected activation field names. Tests cover constructor mappings, layer sequences, pinned outputs, and forward-pass type inference.

Priority: ➖ Normal

Estimated code review effort: 4 (Complex) | ~45 minutes

Change: Refactor

Merge Risk: 🔵 Low · up to 23a9c

Extracted decoders can record the wrong block count for architectures whose encoder and decoder block counts differ. This predates the PR and is minor. The change is otherwise low risk to merge.

Pre-merge checks: 5

Check name	Status	Explanation
Title check		The title clearly summarizes the main M5 changes: shared layer functors, type-parameter refactoring, and one-time MHA head-key construction.
Description check		The description directly explains the layer and architecture refactors, bug fixes, tests, verification results, and known issues covered by the changeset.
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.

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. (35 skipped: 35 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.}

✨ Finishing Touches: 💡 1

🛠️ Fix failing CI checks: 💡

• Commit to this branch
• Create a new PR

[coderabbitai]

Actionable comments posted: 1

🧹 Nitpick comments: 1

test/layers/type_arguments.jl: 1

22-25: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick win

Scan only Julia source files recursively.

The current src/layers/ tree is flat, so this does not cause a current test failure. If a subdirectory is added, readdir(directory; join = true) passes that directory to eachline, and the test can error. A nested .jl file would also be skipped. Use a recursive .jl-file list.

Suggested fix

-    hits = [(file, i)
-            for file in readdir(directory; join = true)
-            for (i, line) in enumerate(eachline(file)) if occursin(r"T::Type\b", line)]
+    files = [joinpath(root, f) for (root, _, fs) in walkdir(directory)
+             for f in fs if endswith(f, ".jl")]
+    hits = [(file, i)
+            for file in files
+            for (i, line) in enumerate(eachline(file)) if occursin(r"T::Type\b", line)]

---

ℹ️ Review info

⚙️ Run configuration

• Configuration used: Organization UI
• Review profile: CHILL
• Plan: Advanced
• Run ID: 76a76ebc-de0e-42e3-8b62-efb07f06060d

📥 Commits

Reviewing files that changed from the base of the PR and between 4c4a832 and 23a9c74.

📒 Files selected for processing: 35

• CHANGELOG.md
• KNOWN_ISSUES.md
• src/architectures/autoencoder.jl
• src/architectures/linear_symplectic_transformer.jl
• src/architectures/standard_transformer_integrator.jl
• src/architectures/symplectic_autoencoder.jl
• src/architectures/symplectic_transformer.jl
• src/architectures/sympnet.jl
• src/architectures/transformer_neural_network.jl
• src/layers/linear_symplectic_attention.jl
• src/layers/multi_head_attention.jl
• src/layers/positional_encoding.jl
• src/layers/psd_like_layer.jl
• src/layers/resnet.jl
• src/layers/symplectic_attention.jl
• src/layers/sympnets.jl
• src/layers/transformer.jl
• src/layers/volume_preserving_attention.jl
• src/layers/volume_preserving_feedforward.jl
• test/architectures/classification_transformer.jl
• test/architectures/forward_pins.jl
• test/architectures/gradient_layer_builders.jl
• test/architectures/linear_symplectic_transformer.jl
• test/architectures/standard_transformer_integrator.jl
• test/architectures/sympnet_integrator.jl
• test/helpers/architecture_pins.jl
• test/integration/documented_defaults.jl
• test/integration/forward_inference.jl
• test/layers/multi_head_attention_allocation.jl
• test/layers/multi_head_attention_key_order.jl
• test/layers/positional_encoding.jl
• test/layers/resnet_bias_encoding.jl
• test/layers/type_arguments.jl
• test/layers/volume_preserving_feedforward.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.

[coderabbitai]

🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win

Pass n_decoder_blocks, not n_encoder_blocks, to the decoder architecture.

UnknownDecoder and UnknownSymplecticDecoder declare the field n_decoder_blocks. This line fills that field from nn.architecture.n_encoder_blocks. The block counts can differ. One example is SymplecticAutoencoder, which has 2 encoder blocks and 3 decoder blocks by default. For an autoencoder with different counts, the extracted decoder's n_decoder_blocks is wrong. The generic method now also handles every SymplecticCompression (for example PSDArch), so this wrong value reaches those architectures too.

Proposed fix

--- "a/src/architectures/autoencoder.jl"
+++ "b/src/architectures/autoencoder.jl"
@@ -233,7 +233,7 @@
 function decoder(nn::NeuralNetwork{<:AutoEncoder})
     DecoderType = last(_unknown_encoder_decoder(nn.architecture))
     NeuralNetwork(
-        DecoderType(nn.architecture.full_dim, nn.architecture.reduced_dim, nn.architecture.n_encoder_blocks),
+        DecoderType(nn.architecture.full_dim, nn.architecture.reduced_dim, nn.architecture.n_decoder_blocks),
         decoder_model(nn.architecture),
         decoder_parameters(nn),
         networkbackend(nn))

📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change

	DecoderType(nn.architecture.full_dim, nn.architecture.reduced_dim, nn.architecture.n_encoder_blocks),
	DecoderType(nn.architecture.full_dim, nn.architecture.reduced_dim, nn.architecture.n_decoder_blocks),

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 93.45%. Comparing base (4c4a832) to head (23a9c74).

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #350      +/-   ##
==========================================
+ Coverage   86.72%   93.45%   +6.72%     
==========================================
  Files          58       58              
  Lines        1341     1222     -119     
==========================================
- Hits         1163     1142      -21     
+ Misses        178       80      -98     

☔ 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]

Closed in favour of #351, which is now the one build of part M5.

#351 takes this PR's ClassificationTransformer(dl) delegation and its keyword test, so all three DataLoader constructors delegate. It keeps its own attention: the map over a tuple of heads here infers Any at 32 heads or more, the default of StandardTransformerIntegrator(32), while #351 maps over 1:n_heads with a key tuple built once per head count.

The branch simplify/part-m5-2 stays on the remote.