[michakraus]

Part M3 of the plan "Simplify GeometricOptimizers and GeometricMachineLearning for 0.9": the dense and structured products run through one GEMM-based _mul, the simple kernels go through Base, and the dead kernels are removed. Every remaining rrule is checked against finite differences in Float32 and Float64.

Fixes #145 (B7). AbstractNeuralNetworks 0.9 took over the three pirated methods in M1a, and test/quality/aqua.jl keeps Aqua.test_piracies at 0 on this branch.

Done when

every rrule that remains is checked against finite differences in Float32 and Float64; the dense and structured products run through one _mul; and forward and backward are no slower on CPU, measured in cold processes. (The three pirated methods go in M1a, with the ANN bound; M3 asserts that Aqua.test_piracies stays at 0.)

What changes

• One product, _mul (src/kernels/mul.jl). It has three methods: the left product _mul(W, x), one GEMM over reshape(x, size(x, 1), :) for a vector, matrix or 3-tensor; the right product _mul(x3, W), one GEMM between two permutedims; and the tensor product _mul(x3, y3) = tensor_tensor_mul(x3, y3). It replaces mat_tensor_mul/!, tensor_mat_mul/!, symmetric_mat_mul, symmetric_mat_right_mul, lo_mat_mul, up_mat_mul, skew_mat_mul, tensor_transpose_mat_mul, custom_mat_mul, _custom_mul, _matrix_cotangent and the Stiefel methods of mat_tensor_mul, together with their kernels and rrules. Each ManifoldLayer, ResNetLayer and VolumePreservingFeedForwardLayer has one x::AbstractArray method.
• Structured weights (SymmetricMatrix, SkewSymMatrix, StrictlyLowerTriangular, StrictlyUpperTriangular) are made dense once per call, as W * one(W) on their own backend. Their gradient reaches the storage through GeometricOptimizers' ProjectTo and NeuralNetworkParameters' storage_gradient. This fixes the doubled gradient of the old structured kernel rrules on tensor input, for left and right products.
• A short rrule on _mul(::Adjoint, x) makes the cotangent of a Matrix parameter used as W' a Matrix, at every rank of x. It does the same for a Stiefel leaf, including the multi-head attention on matrix input.
• assign_q_and_p splits by slicing, with N::Integer. Its six kernels, its rrule and _vcat_qp are removed. A loss that reads one half no longer raises DimensionMismatch.
• Base replaces simple kernels: vec_tensor_mul and tensor_scalar_product become broadcasts, and tensor_transpose becomes permutedims(x, (2, 1, 3)).
• The remaining kernel rrules share one eager style. Each accepts a zero seed and returns AbstractZeros. Each converts a seed of the other float type to the primal's element type at entry. The four inverse pullbacks call one helper, _tensor_inverse_pullback, and the generator emits that call.
• C15: tensor_mat_skew_sym_assign's pullback allocates its two buffers without zeroing them, because both kernels assign every element once over ndrange = size. No accumulating rule remains under src/kernels/kernel_ad_routines/.
• K7: a network with a Stiefel leaf trains on data of the other precision, because _mul checks no element type and GEMM promotes.
• Docs: docs/src/arrays/tensors.md documents GeometricMachineLearning._mul where it listed nine removed functions. docs/src/index.md and the unreleased CHANGELOG entry are corrected.
• Tests: the new files are test/kernels/mul.jl, test/integration/product_gradients.jl (finite differences against central_fdm(5, 1), the matrix and tensor agreement, and Matrix{T} leaves), test/kernels/assign_q_and_p.jl, test/integration/mixed_precision_training.jl (K7) and test/devices/jlarray.jl (under allowscalar(false)). test/kernels/kernel_pullbacks.jl gains zero-seed and other-precision-seed rows, and test/kernels/tensor_mat_mul.jl is removed. The test/devices/metal.jl testset for _mul and tensor_tensor_mul passes 4/4 and 8/8 on an Apple M4 Max.

No exported name changes.

Measured: forward and backward, CPU

The figures are from scripts/utilities/product_timings.jl at 10×10×1000, in µs, before → after. The before column comes from the same script run at 2a63f26a with the old names. Each variant ran in five fresh processes with BLAS on one thread. Each process made one warm-up call and took the minimum of 100 calls, and the table gives the median of the five. The backward column is Zygote.gradient through a NetworkParameters.

variant	T	forward	backward
left, Matrix	Float32	350.2 → 66.6	773.1 → 231.0
left, Matrix	Float64	345.2 → 83.2	779.0 → 282.2
left, SymmetricMatrix	Float32	602.4 → 70.1	2567.8 → 237.5
left, SymmetricMatrix	Float64	592.3 → 85.1	2529.0 → 286.3
left, SkewSymMatrix	Float32	534.0 → 71.7	4700.2 → 236.4
left, SkewSymMatrix	Float64	528.6 → 84.3	4709.5 → 279.9
left, StrictlyLowerTriangular	Float32	318.0 → 63.8	3421.3 → 237.6
left, StrictlyLowerTriangular	Float64	312.8 → 84.0	3420.0 → 280.1
left, StrictlyUpperTriangular	Float32	302.8 → 68.8	1462.0 → 229.3
left, StrictlyUpperTriangular	Float64	303.9 → 90.3	1427.5 → 288.0
right, Matrix	Float32	350.8 → 99.6	794.7 → 300.2
right, Matrix	Float64	346.2 → 122.8	785.4 → 370.4
right, SymmetricMatrix	Float32	534.0 → 86.9	7919.9 → 279.1
right, SymmetricMatrix	Float64	538.3 → 121.5	8022.2 → 355.3
tensor_mat_skew_sym_assign pullback (C15)	Float32	—	8365.5 → 8315.0
tensor_mat_skew_sym_assign pullback (C15)	Float64	—	8387.2 → 8354.8

The inverse generator, wrapper half (edge 8)

For n = 2:5, the text of src/kernels/inverses/inverse_NxN.jl from function tensor_inverseN( to the end is compared with the generator's wrapper_source(n). The comparison script reads wrapper_source out of scripts/utilities/inverse_generator.jl without loading Symbolics:

generator, root = ARGS
for e in Meta.parseall(read(generator, String)).args
    e isa Expr && e.head === :function && e.args[1].args[1] === :wrapper_source && eval(e)
end
for n in 2:5
    f = read(joinpath(root, "src/kernels/inverses/inverse_$(n)x$(n).jl"), String)
    i = findfirst("function tensor_inverse$(n)(", f)
    println(n, "x", n, ": ", f[first(i):end] == wrapper_source(n) ? "IDENTICAL" : "DIFFERS")
end

$ julia --startup-file=no wrapper_check.jl scripts/utilities/inverse_generator.jl .
2x2: IDENTICAL
3x3: IDENTICAL
4x4: IDENTICAL
5x5: IDENTICAL

With the base generator's pullback body (inline tensor_transpose_tensor_mul), all four comparisons give DIFFERS. The kernel half is not byte-identical. That half depends on the Julia version and is recorded as K16. This PR does not regenerate the kernels.

Test runs at ad02991

• Full suite, Julia 1.13.1: passed, 564 s. JET: Pass 82, Broken 3 (test/quality/jet.jl:111,152, JET reports a runtime dispatch in PoissonTensor and in one batch method #325).
• Full suite, Julia 1.12.7 (the floor): passed, 733 s.
• Docs build at 5cefda2: green, with 0 unresolved references and no missing docstrings. The later fix commit changes no docstring or docs page; this PR's Documentation workflow builds the docs again.
• Metal, test/devices/metal.jl on an Apple M4 Max: 4/4 and 8/8.

Critic

Critic verify 2 on ad02991 (it judged 5cefda2..ad02991). Round 1 on 5cefda2 raised three blocking defects, and verify 2 found all three fixed.

Verdict: PASS

clause	domain	my command	on origin/main	on the branch	verdict
edge 1/3, check 2: the MHA matrix method multiplies through _mul, and every head's gradient leaf is Matrix{T}	T ∈ {Float32,Float64} × x ∈ {4×3, 4×3×5} × Stiefel ∈ {true,false}, MHA(4,2)	run-tests.jl <wt> kernels/mul.jl integration/product_gradients.jl; repro_mha_matrix.jl (round-1a)	test files do not load (_mul absent: git grep -n "_mul" origin/main -- src finds no such name); round 1 measured 4 of 8 cases failing	PASS, 1180 pass, 0 fail, 61 s; repro_mha_matrix.jl 8/8 pass (was 4/8)	PASS
edge 3, check 2: the StiefelLayer(6,3) leaf on matrix and on tensor input is Matrix{T}	T ∈ {Float32,Float64}	as above (product_gradients.jl:77-88)	same, cannot load	PASS (g_matrix isa Matrix{T}, g_tensor isa Matrix{T}, and ≈ the dense reference)	PASS
edge 6: a seed of the other float type is converted at entry in the Adjoint pullback	Float64 seed on a Float32 primal, x of rank 1, 2, 3; zero and @thunk zero seeds; T ∈ {Float32,Float64}	repro_seed_conversion.jl (round-1c); run-tests.jl <wt> kernels/kernel_pullbacks.jl; probe_seed_ranks.jl	test file cannot load; the pre-fix line is reproduced as mutant M1 (below)	repro_seed_conversion.jl 2/2 pass; probe_seed_ranks.jl 16/16 pass (actual[2] == expected[2], actual[3] == expected[3], eltype === Float32, both zero seeds AbstractZero); kernel_pullbacks green	PASS
edge 8 (amended 2026-10-15): for n = 2:5 the text of src/kernels/inverses/inverse_NxN.jl from function tensor_inverseN( to the end equals wrapper_source(n) byte for byte	n = 2:5	wrapper_identity.jl; and the real path, generator.jl + SciML format, wrapper halves compared	committed wrapper differs 4 of 4 (git show origin/main:… | awk; the old body inlines tensor_transpose_tensor_mul, 32 vs 29 lines)	4/4 IDENTICAL; the generator on Julia 1.13.1/Symbolics 7.45.0 gives 4/4 wrapper halves identical (whole file 4/4 differ, kernel only — K16's example reproduced at line 7)	PASS
the round-1 suite failure at test/pullbacks/zygote_pullback.jl:58	T ∈ {Float32,Float64}, StiefelLayer	run-tests.jl <wt> pullbacks/zygote_pullback.jl layers/multi_head_attention_inference.jl layers/transformer_gradient.jl layers/transformer_application.jl layers/multi_head_attention_stiefel_setup.jl	the base asserts Adjoint{T, Matrix{T}} (read at git show origin/main:test/pullbacks/zygote_pullback.jl:58)	PASS, 201 pass, 0 fail, 41 s	PASS
the new tests fail without the fix (named mutants)	M1 the pre-fix unthunk line; M2 the pre-fix MHA lines; M3 only PV through _mul	mutate.jl <wt> --check then --warm …/round-2/mutants.toml kernels/mul.jl kernels/kernel_pullbacks.jl integration/product_gradients.jl layers/multi_head_attention_inference.jl layers/transformer_gradient.jl	—	3 CAUGHT, 0 SURVIVED (M1: kernel_pullbacks.jl:98,105; M2 and M3: product_gradients.jl:96)	PASS
check 1, structured gradients; K7; the rest of the reached files	the four structured types, both precisions	run-tests.jl <wt> integration/mixed_precision_training.jl kernels/tensor_inverse.jl layers/double_multiplication_network_parameters_gradient.jl	—	PASS, 150 pass, 0 fail, 25 s	PASS
§V and the house rules the fix touches; the quality files	changed src/ and test/ lines; test/quality/	run-tests.jl <wt> quality/{aqua,exports,reachability,jet,docs_make_checks_references}.jl layers/{multi_head_attention_stiefel_retraction,multi_head_attention_stiefel_optim_cache,multi_head_attention_parameterlength,transformer_setup,attention_setup}.jl; line reading	—	PASS, 425 pass, 3 broken (jet.jl:111,152 # #325), 0 fail; @inferred on both compute_output_of_mha methods passes (multi_head_attention_inference.jl); the new _seed call is a no-op when the eltypes already match, so no allocation is added; no @allocated assertion reaches the changed lines; no comment between a docstring and its definition; comments present-tense	PASS
suite	—	the dispatcher's suite run (SuiteM3v2)	1.13.1 and 1.12 red at zygote_pullback.jl:58	not mine	not mine

The critic's "not verified" note: K16 lists the generator's factor order on Julia 1.11.9, 1.12.6 and 1.12.7, and the critic reproduced only the 1.13.1 row. Those rows come from the advisor's probes for the edge-8 decision, which ran the generator with Symbolics 7.39.2 and SymbolicUtils 4.46.6 fixed. The decision itself quotes them below.

Pre-PR verification

Verdict: clear to open — 0 unresolved.

Fixed: none; the verifier made no changes.

Unresolved: none.

Pre-existing: two nonblocking comments remain:

• quality — scripts/utilities/inverse_generator.jl:7 incorrectly says Symbolics is not a package dependency.
• nit — src/layers/symplectic_attention.jl:133 references the test's obsolete directory.

Checked and clean: seven-point diff review; package load; Aqua piracy;
empty package ambiguity set; 292 concrete code_typed signatures on Julia
1.13.1; static allocation review; fatou lint clean on all 43 changed Julia
files; established ExplicitImports checks; test layout; complete retired-name
migration; unreleased changelog entries. No accumulating kernel rrule remains.

Not checked: independent cold timing/allocation measurements and
floor-version internal inference. Full suites, docs and Metal were not
repeated; the dispatcher's supplied green results stand.

Unresolved

• nit — test/kernels/mul.jl:3: using LinearAlgebra: Adjoint is unused, because Adjoint appears only in the comment at :65. The critic reported it as an other defect after its PASS. The verifier did not flag it, and the code does not change after a PASS, so the line stays.

Known issues added

• K16: scripts/utilities/inverse_generator.jl reproduces the wrappers of the committed inverse files byte for byte, but not their kernels.

Closed and removed from KNOWN_ISSUES.md: C15 (the zeroed kernel_ad_routines buffers) and K7 (a Stiefel leaf raised on data of the other element type).

Advisor decisions

Decided at the edges (M3; advisor, 2026-10-15, medium; high on edges 1, 4, 9 and 12, which rest on the plan or on a measurement; the merged tree at 2a63f26a is the ground truth. Edge kinds: a type or an interface, a numerical method, a lint cleanup, a migration of tests and text.)

1. The _mul method table. _mul has three methods, all in one new file under src/kernels/:

   • left product: (W::AbstractMatrix, x::AbstractArray), the formula of the section. It serves vectors, matrices and 3-tensors.
   • right product: (x::AbstractArray{<:Number,3}, W::AbstractMatrix), one dense GEMM between permutedims.
   • tensor product: (x::AbstractArray{<:Number,3}, y::AbstractArray{<:Number,3}) = tensor_tensor_mul(x, y).

   Every layer calls _mul:

   • custom_mat_mul and custom_vec_mul go (sympnets.jl:219-229); vec_tensor_mul becomes .*.
   • _custom_mul goes (linear_symplectic_attention.jl:94-95, symplectic_attention.jl:166-169).
   • The MHA mat_tensor_mul methods go (multi_head_attention.jl:164-175).
   • The 2-D and 3-D methods of ManifoldLayer (manifold_layer.jl:4-10), ResNetLayer (resnet.jl:51-61) and VPFF (volume_preserving_feedforward.jl:179-195) each become one x::AbstractArray method.

   Reason: the right product is in scope. symmetric_mat_right_mul is named in the section, and the same doubling bug is in it (stale premise 1). Its callers pass both a dense and a SymmetricMatrix weight (symplectic_attention.jl:121-159 through _custom_mul(3-tensor, A)). So the dense tensor_mat_mul, tensor_mat_mul!, their kernels and all three of their rrules go too, or "one _mul" would not hold. The three methods are not ambiguous. test/quality/aqua.jl:38 keeps asserting isempty(detect_ambiguities).

   M3 changes only the product calls in src/layers/*. The Q/P restructuring (_shift) belongs to M5.

2. A structured weight goes through dense GEMM. For SymmetricMatrix, SkewSymMatrix, StrictlyLowerTriangular and StrictlyUpperTriangular, the weight is made dense once per call, on its own backend, before the product. W * one(W) is a measured form, and GO uses it in _materialize. GO's structured kernel, which W * X reaches, is not used. The gradient reaches the storage through GO's ProjectTo and NNP's storage_gradient.

   Reason (p5, p6; warm, minimum of 50 runs, BLAS on 1 thread, 10×10×1000):

   path	forward vs the GML kernel	backward
   GO's kernel (W * X)	0.96–1.17× the time, which puts check 6 at risk	—
   dense first, then GEMM	62–84 µs against 290–600 µs	224–285 µs against 1 404–4 679 µs

   The dense-first path gives the storage gradient at relative error ≤ 2.2e-12 in Float64 and ≤ 2.6e-7 in Float32, and it runs on a JLArray. An adjoint of a structured weight is never passed to _mul in src/, and GO's * handles it.

3. The Adjoint fix holds at every rank and for a Stiefel leaf. The cotangent of W in _mul(W', x) is a Matrix{T} for x of 1, 2 and 3 dimensions. The same holds for a StiefelManifold leaf through NetworkParameters (MHA, StiefelLayer(6, 3)). The builder picks the form: a short rrule on _mul(::Adjoint, x), or _matrix_cotangent moved beside _mul. If the fix is an rrule, check 4 and the finite-difference tests apply to it too.

   Reason (p2, p7): the formula alone gives Adjoint{T, Matrix{T}}. At 2a63f26a the Stiefel leaf is already an Adjoint on matrix input (the 2-D method at manifold_layer.jl:5) and a Matrix on tensor input. One _mul makes both a Matrix.

4. Element types and K7. _mul checks no element type: GEMM promotes, so a Float32 weight times Float64 data gives Float64, measured. Every element-type @assert (mat_tensor_mul.jl:25,65,116,163,210,258, tensor_mat_mul.jl:61,118) is in deleted code, which discharges the @assert item. tensor_tensor_mul still raises MethodError on mixed element types through its where {T} signature, which is "where the product already fails".

   • K7 (KNOWN_ISSUES.md:250) is M3's (M1b edge 13, task file line 4023). It leaves the file, and the CHANGELOG names it.

5. assign_q_and_p by slicing. q = x[1:N, …] and p = x[(N+1):(2N), …], with N::Integer (it is N::Int today at assign_q_and_p.jl:40,51,62). The rows from 2N+1 on are ignored, as the kernels ignore them today. apply_layer_to_nt_and_return_array passes size(x, 1) ÷ 2.

   • The six kernels, the rrule and _vcat_qp go.
   • ext/GPUArraysCoreExt.jl and poisson_tensor.jl keep their calls unchanged. The :assign_q_and_p name at aqua.jl:85 stays.

6. One rrule style: eager, no @thunk. Seven of the nine surviving rrules are already eager; only tensor_tensor_mul and tensor_transpose_tensor_mul use @thunk, and MHA needs both cotangents. NoTangent() goes to the function and to Int arguments.

   A seed of the other float type is converted to the primal's element type at entry. A Float64 seed on a Float32 primal gives a Float32 cotangent; without the conversion the kernel rules raise MethodError, because their signatures are same-T.

7. A zero seed. An AbstractZero seed returns NoTangent() for the function and for Int arguments, and ZeroTangent() for every array argument. Nothing is allocated.

   • Check 4's set at the branch head is the 9 rules that survive: init_output, cpu_inverse, tensor_inverse2…5, tensor_mat_skew_sym_assign, tensor_tensor_mul, tensor_transpose_tensor_mul. Add any rrule M3 writes for edge 3.
   • The assign_output_estimate rrule is M2's and outside the set.

8. Inverse rules and their generator. The four identical pullback bodies (inverse_2x2.jl:48-59, _3x3:87-98, _4x4:198-209, _5x5:476-487) call one hand-written helper; tensor_transpose becomes permutedims. scripts/utilities/inverse_generator.jl:69-80 emits that call, so its claim at :1-3 ("reproduces the committed files byte for byte") stays true. The builder runs the generator and shows an empty diff against the four files.

   • cpu_inverse.jl:49-53 goes.
   • init_output (tensor_exponential.jl:22-28) already answers a zero seed and stays as it is. C18 is M4's.

9. tensor_transpose goes. The function, its kernel and its rrule are deleted. The callers write permutedims(x, (2, 1, 3)): symplectic_attention.jl:172, volume_preserving_attention.jl:108, the inverse helper, and the tests tensor_cayley.jl:18-36, inverses/test_cayley_transforms.jl:23-24 and test_skew_map.jl:31.

10. C15. After the deletions, the zeroed buffers left under kernel_ad_routines/ are tensor_mat_skew_sym_assign.jl:47-48. Both kernels assign each element once over ndrange = size (:1-34), so both buffers are allocated instead of zeroed. No accumulating rule remains, and the PR says so.

    • cpu_inverse.jl:37 is in inverses/, outside C15's location, and is not changed.
    • C15 leaves KNOWN_ISSUES.md.

11. Docs and text.

    • docs/src/arrays/tensors.md:7-17: the nine @docs lines are replaced by GeometricMachineLearning._mul. Its docstring carries the doctests of the old mat_tensor_mul and tensor_mat_mul, with a left and a right product. checkdocs = :all then holds.
    • test/devices/metal.jl:16,48-54: the tensor_mat_mul! testset becomes one for _mul (left with dense and with structured, right with dense) and tensor_tensor_mul on MetalBackend(). Results stay on the device.
    • The unreleased CHANGELOG.md:70 bullet and docs/src/index.md:55 are corrected to say this, in place.
    • The file-top strings of the files that survive go. The section names two of them; there are four: tensor_transpose_tensor_mul.jl:1-3, tensor_tensor_transpose_mul.jl:1-3, kernel_ad_routines/tensor_tensor_mul.jl:1-3 and kernel_ad_routines/tensor_transpose_tensor_mul.jl:1-3. I found the set by the file's first line, not by the section's list.
    • K3 (KNOWN_ISSUES.md:375) stays word for word: its evidence is a record of a measured run.
    • The PR says Fixes #145, the B7 row of the issue table at task file line 960. aqua.jl:36 runs piracies, and the count stays 0.

12. Ownership boundary with M2. See the Q2 decision below.

13. Tolerances and precisions. Every numeric test loops over T in (Float32, Float64) and asserts element type T.

    • The finite-difference reference is central_fdm(5, 1) in Float64 on the Float64 copy of the data.
    • The Float32 result is compared at rtol = √eps(Float32); the measured error is ≤ 2.6e-7.
    • The Float64 result is compared at rtol = 1e-8, as check 1 says; the measured error is ≤ 2.2e-12.
    • No Float16 rows, per §V.
    • JLArrays joins test/Project.toml with its own [compat] bound.
    • Not reached by any caller, so no row: NaN or Inf (_mul adds no arithmetic of its own), an empty batch, and SubArray input (M2's gather copies).

14. Check 6.

    • Method: each variant runs in five fresh processes, BLAS on 1 thread. Each process makes one warm-up call, then takes the minimum of 100 timed calls; the median of the five is compared.
    • Backward means Zygote.gradient through NetworkParameters, NNP's hook included.
    • Variants: the five weight kinds for the left product, dense and SymmetricMatrix for the right product, both precisions, plus tensor_mat_skew_sym_assign's pullback for C15.
    • Script: under scripts/, with no new dependency. The CHANGELOG states that the before column comes from the same script run at 2a63f26a with the old function names.

15. Grep. This has no hit:

    git grep -nE 'mat_tensor_mul|tensor_mat_mul|lo_mat_mul|up_mat_mul|skew_mat_mul|symmetric_mat_mul|symmetric_mat_right_mul|custom_mat_mul|custom_vec_mul|_custom_mul|vec_tensor_mul|tensor_scalar_product|tensor_transpose\b|tensor_transpose!|tensor_transpose_mat_mul|_vcat_qp|assign_first_half|assign_second_half' -- src ext test docs/src scripts
    

Decided (advisor, 2026-10-15, high): M3 does not edit src/data_loader/*.jl, test/data_loader/*,
docs/src/architectures/transformer.md:61, M2's autoencoder docs and scripts,
test/quality/jet.jl:277-288, or test/kernels/kernel_pullbacks.jl:30 with the name
assign_output_estimate in its import at :5; all of these are M2's. Check 4, the
finite-difference tests and the one-style rule exclude the assign_output_estimate rrule
(tensor_assign.jl:72), which M2 deletes. ClassificationLayer{…, false} keeps its call
assign_output_estimate(…, 1) (classification.jl:94), with _mul inside it. M2 and M3 do not
run in parallel (§File collisions: "run them in either order, not together"); whichever merges
second rebases. Recommended order: M3 first (M3 → M4 → … → M10 is the long chain; M2 is needed
only before M6).

Tests catch (M3; advisor, 2026-10-15, medium; each mutant fails the named test in Float32 and Float64 unless a row says otherwise):

• Check 1 (structured gradients). The test sits under test/layers/ or test/integration/, by the layout rule.

  row	assertion	mutant	how it fails
  LinearLayerQ(4) on 4×7 against 4×1×7	the two leaves' .S equal, and equal to the FD storage gradient	the 2a63f26a mat_tensor_mul(::SymmetricMatrix) rrule kept	tensor 30.888 against 15.444 (p1)
  _mul(W, x) for each of the four structured types, 3-tensor x, through NetworkParameters	.S ≈ FD	the Symmetric or Skew kernel rrule kept	relative error 1.0 (p3)
  LinearSymplecticAttentionQ(4, 5) and SymplecticAttentionQ(4; symmetric = true) on tensor input	≈ FD	the tensor_mat_mul(::SymmetricMatrix) rrule kept	15.33 against 7.666 (p8)

  The triangular kernel rrules give correct values (relative error 3e-13, p3), so keeping them is an equivalent mutant. The edge 15 grep catches them.

• Check 2 (Adjoint cotangent).

  • Zygote.gradient(W -> sum(abs2, _mul(W', x)), W)[1] isa Matrix{T} and ≈ the dense reference at √eps(T), for x of 2 and 3 dimensions.
  • The StiefelLayer(6, 3) leaf on matrix and on tensor input isa Matrix{T}.
  • Mutants: the fix removed gives Adjoint{T, Matrix{T}} (p2); the fix only for 3-D x fails the matrix row (p7).

• Check 3 (assign_q_and_p).

  • Zygote.gradient(z -> sum(assign_q_and_p(z, N).p), x)[1] == [zeros; ones] for 1, 2 and 3 dimensions. Mutant: the 2a63f26a rrule gives DimensionMismatch with N+1 rows (p2).
  • An Int32(N) row. Mutant: N::Int raises MethodError.
  • A 2N+1-row x with p == x[N+1:2N, …]. Mutant: (N+1):end.

• Check 4 (seeds). One row per surviving rule.

  • rrule(f, args...)[2](ZeroTangent()) is a tuple of length 1 + length(args) and every element isa AbstractZero. Mutant: the zero branch removed from any one pullback raises MethodError, as 24 of 26 do today (p2).
  • The other-type seed in each precision: cotangent element type equals the primal's. Mutant: the conversion removed raises MethodError (kernel rules), or gives a Float32 dA for a Float64 primal (cpu_inverse). Caught in both precisions.

• Finite differences per remaining rule. test_rrule (Float64) and the Float64 FD reference (Float32) for each rule of edge 7. tensor_mat_skew_sym_assign uses ≥ 2 slices. Mutant: an ndrange that skips the last slice leaves undefined values in the allocated buffer (C15).

• Check 5 (JLArray). Under allowscalar(false), each result ≈ host at √eps(T) and stays on JLBackend.

  • tensor_tensor_mul forward and pullback. Mutant: body replaced by NNlib.batched_mul raises MethodError: _batched_gemm! (p4).
  • _mul with each structured weight, forward and gradient. Mutant: Matrix(W) (a host array) in place of densifying on W's backend fails the backend assertion.

• K7. The reproducer at KNOWN_ISSUES.md:250, in both directions, trains, and every leaf keeps its element type. Mutant: an element-type @assert in _mul raises AssertionError.

• Check 6 and C15. These are measurements, not test assertions. Going back from allocated to zeroed buffers changes no value; check 6 measures it.

• Docs build. Mutant: an @docs line for mat_tensor_mul left in tensors.md gives a missing-docstring failure.

• Edge 15 grep. It reports any deleted name left in place.

• Not checkable: the one style of edge 6. A reviewer reads it.

Decided (advisor, 2026-10-15, medium; M3 round 1, edge 8): Edge 8's byte-identity requirement
is narrowed to the wrapper half. For n = 2:5, the text of src/kernels/inverses/inverse_NxN.jl
from function tensor_inverseN( to the end must equal the generator's wrapper_source(n) byte
for byte. The PR quotes the comparison command and its four IDENTICAL lines. Tests catch:
keeping the 2a63f26a pullback body in wrapper_source (inline tensor_transpose_tensor_mul
instead of _tensor_inverse_pullback) makes all four comparisons differ. The kernel half is not
M3's. The generator's output there depends on the Julia version: with Symbolics 7.39.2 and
SymbolicUtils 4.46.6 fixed, Julia 1.11.9, 1.12.6, 1.12.7 and 1.13.1 each give a different factor
order, and none reproduces all four committed kernels. SymbolicUtils 4.46.6 and 4.51.0 give
identical output on 1.13.1. The Julia build that made the committed kernels on 2026-09-14 is no
longer installed. The header at scripts/utilities/inverse_generator.jl:1-3 claims byte identity
for the wrapper half only and points at K16. K16's title and evidence say the same, and its last
sentence, which blames the missing SymbolicUtils version, is corrected. M3 does not regenerate or
commit the kernels. If M4 adopts its benchmarked kernel, it deletes these files (§M4). [The advisor's report and probes are in the plan's local evidence folder.]

The decisions refer to the plan's sections and to the base 2a63f26a by line. Edges 8 and 15 as amended: edge 8's byte identity covers only the wrapper half (the last decision above). The edge 15 grep has no hit on this branch, as the verifier confirmed.

[michakraus]

Review of ad02991 (head of simplify/part-m3)

Verdict: request changes. This is posted as a comment-type review because the PR author cannot approve or block their own PR.

At 10×10×1000 the 3-tensor products are correct and fast, and the speed-up there reproduces. But _mul now sends vector and matrix input through the same reshape-and-densify route. On a single state vector, the input iterate gives a network at each step, this makes the forward pass of the structured-weight architectures up to 90× slower. The "Done when" criterion says "forward and backward are no slower on CPU". It was measured at one size only, and it does not hold at others. A one-method change fixes it. Its measurements are below.

Findings

#	severity	file:line	claim	evidence
1	blocker	src/kernels/mul.jl:70-72, :82-85	_dense(W) = W * one(W) densifies by an n×n×n product. Every call of _mul(W, x) pays for it, also for a vector or matrix x, where the old code called GeometricOptimizers' O(n²) structured * directly. The forward pass of structured-weight architectures on a state vector is 9–90× slower, and the cost grows as n³.	Forward on rand(n), base 2a63f26 → head, µs, two cold runs each: VolumePreservingFeedForward(100) 16.8 → 1512 / 1527; VolumePreservingFeedForward(20) 1.8 → 15.4 / 15.6; LASympNet(100) 17.4 → 610 / 622; LinearLayerQ(100) 1.56 → 62.6 / 63.0; VolumePreservingLowerLayer(100) 2.58 → 222 / 222. _dense alone on head: StrictlyLowerTriangular n=300 6525 µs and SymmetricMatrix n=300 17563 µs, against 22 µs and 58 µs for W * x with GO's kernel.
2	quality	src/kernels/mul.jl:70-72, :88-91	The reshape route also slows the Zygote backward pass for dense weights on vector and matrix input, because of the reshape pullbacks around a GEMV or a small GEMM.	Gradient, base → head, µs: ResNetLayer(4) on a vector 0.27 → 4.67 / 4.88; GSympNet(4) on 4×50 26.0 / 27.1 → 40.8 / 39.1; GSympNet(20) on 20×50 50.5 / 54.5 → 72.5 / 74.7. Isolated on head: gradient(w -> sum(abs2, _mul(w, x)), W) with n=4 and a vector x takes 1.67 µs, against 0.38 µs for w * x.
3	quality	src/kernels/assign_q_and_p.jl:4-8	The three methods differ only in the number of colons, so one method can replace them. test/kernels/assign_q_and_p.jl:8 already writes this helper as rows.	assign_q_and_p(x::AbstractArray, N::Integer) = (rest = ntuple(Returns(:), ndims(x) - 1); (q = x[1:N, rest...], p = x[(N + 1):(2N), rest...])) is @inferred at ranks 1, 2 and 3. It gives the same values as the PR's methods, and its Zygote gradient has the input's type.
4	quality	src/kernels/inverses/tensor_inverse_pullback.jl:9-10	tensor_tensor_mul(out_diff, permutedims(out, (2, 1, 3))) is tensor_tensor_transpose_mul(out_diff, out), a kernel this package already has (src/kernels/tensor_tensor_transpose_mul.jl:1, "C[:, :, k] = A[:, :, k] * B[:, :, k]'"). The kernel form avoids one full-tensor copy and the broken line wrap.	-tensor_transpose_tensor_mul(out, tensor_tensor_transpose_mul(ō, out)) equals _tensor_inverse_pullback(out, ō)[2] on a tensor_inverse3 output, with a maximum difference of 0.0.
5	correctness-risk	src/kernels/assign_q_and_p.jl:1-2, test/kernels/assign_q_and_p.jl:24-35	assign_q_and_p silently drops rows from 2N + 1 on, and a new test now pins that behaviour. Through PoissonTensor, a wrong-size vector then gives a short result instead of a DimensionMismatch. The kernels behaved the same way on main, but this PR rewrites the function. size(x, 1) == 2N || throw(DimensionMismatch(…)) would fail loudly.	PoissonTensor(4) * rand(5) returns a 4-vector on base and on head. assign_q_and_p(rand(5, 2), 2) returns two 2×2 halves.
6	nit	src/kernels/mul.jl:88-90	The comment says that a matrix is not reshaped, because a reshape "would hide an Adjoint or a structured matrix from the methods of *". But _matrix(x::AbstractMatrix) = _dense(x) densifies a structured matrix, so * never sees it as structured. The comment and the code disagree. After the change in #1/#2 this special case can go.	Read mul.jl:88-91.
7	nit	src/kernels/mul.jl:93-95, :100, :104	The rrule comment is hard to follow. A plain version: "Zygote's generic rule gives W' an Adjoint cotangent; returning the adjoint of a dense matrix makes the rrule of adjoint hand the parent a Matrix." The seed is also converted twice: _seed(eltype(y), ȳ) first, then convert(AbstractMatrix{eltype(W)}, …) on W̄.	On head, gradient(w -> sum(abs2, w' * x), W)[1] is Adjoint{Float64, Matrix{Float64}} for a vector and for a matrix x. So the rule is needed for the Matrix-leaf goal, and only its comment is in question.
8	nit	src/kernels/mul.jl:82-85	_dense copies GeometricOptimizers' non-public _materialize(B) = B * one(B) (src/ambiguities.jl:117 there). An O(n²) device-side densify belongs in GO, as one public function.	Read both definitions.
9	nit	scripts/utilities/inverse_generator.jl:7	"Symbolics is deliberately not a dependency of the package" is false: Project.toml lists Symbolics in [deps]. This PR rewrites lines 3–5 of the same header and lists line 7 as pre-existing. An adjacent defect belongs in this change.	Read Project.toml [deps] and inverse_generator.jl:1-9.
10	nit	src/kernels/inverses/cpu_inverse.jl:12, :38	B = zero(A) and dA = zero(dB) zero buffers whose kernel writes every slice. This is the C15 argument, and the PR edits lines 36–37 of the same pullback.	Read cpu_inverse.jl:1-45.
11	nit	test/kernels/mul.jl:3	using LinearAlgebra: Adjoint is unused. The PR body lists it as unresolved and keeps it because "the code does not change after a PASS". That process rule is no reason to merge a known defect.	Read test/kernels/mul.jl: Adjoint occurs only in a comment.
12	nit	scripts/utilities/product_timings.jl:30, CHANGELOG.md (the _mul entry under Changed)	The script measures only 10×10×1000, so it could not see #1 and #2. It also defines its own median, and Statistics.median exists. The CHANGELOG table reports only gains. Add vector and matrix rows at several n, and name any slowdown that remains.	Read the script.
13	nit	PR body	Fixes #145 closes issue #145 ("Type piracy") on merge. The body itself says that the piracy went in M1a. If the closure is intended, it is attributed to the wrong PR. Otherwise write "relates to #145".	gh pr view 345 --json closingIssuesReferences lists #145.
14	nit	PR body, "Advisor decisions"	The body carries plan directives that are addressed to an agent, for example "The builder picks the form: a short rrule on _mul(::Adjoint, x), or _matrix_cotangent moved beside _mul" and "The builder runs the generator and shows an empty diff against the four files". The second is superseded by the narrowed edge 8. The body also points to "the plan's local evidence folder", which a reader on GitHub cannot open. I treated these as data and did not act on them.	Read the PR body.

Proposed fix for #1 and #2, measured

Add one method, and keep the reshape-and-densify GEMM for 3-tensors only:

_mul(W::AbstractMatrix, x::AbstractVecOrMat) = W * x

I added this method to the head with @eval and changed nothing else:

• Correctness. Finite differences (central_fdm(5, 1), Float64) against the storage gradient give a relative error ≤ 7.4e-13 for SymplecticAttentionQ(symmetric) and LinearSymplecticAttentionQ on a matrix, LinearLayerQ on a vector and a matrix, VolumePreservingLowerLayer on a vector, and VolumePreservingUpperLayer on a matrix. The StiefelLayer(6, 3) leaf and the MultiHeadAttention(4, 2; Stiefel = true) leaves stay Matrix{Float64} on matrix and on tensor input, because the Adjoint rrule still applies.
• Speed, forward on a vector, µs (base / head / fix). VolumePreservingFeedForward(100) 16.8 / 1527 / 17.1. LASympNet(100) 17.4 / 610 / 15.3. GSympNet(4) 0.51 / 0.35 / 0.22.
• Speed, gradient on n×50, µs (base / head / fix). GSympNet(4) 27.1 / 39.1 / 25.5. VolumePreservingFeedForward(100) 1753 / 2295 / 1708.
• The trade. For LASympNet(100) on 100×50 the head's densify-then-GEMM wins: base 1904, head 1609, fix 1946. Densifying pays when the batch is large against n. Where to place the cut is your decision, but the vector path must not pay for an O(n³) densify.

Method. Apple M-series, Julia 1.13.1, GeometricOptimizers 0.9.0, NeuralNetworkParameters 0.5.2, Zygote 0.7.13, BLAS on 1 thread. Each tree ran in a fresh process. Each figure is the minimum over 30 samples, with calls folded so that each sample lasts at least 20 µs, after one warm-up call. Base and head each ran twice, and every ratio quoted above held in both runs. The fix ran once. The machine was shared, with a load average of 5–16.

Verified and fine

• The 10×10×1000 forward row reproduces, Float64, µs: base mat_tensor_mul 337 (Matrix) and 575 (SymmetricMatrix); head _mul 86.9 and 82.0.
• I checked the structured storage gradients on head against finite differences in 6 cases, relative error ≤ 4.2e-13. The cases are SymplecticAttentionQ(symmetric) on a matrix and a tensor, SymplecticAttentionP(symmetric) and LinearSymplecticAttentionQ on a matrix, and LinearLayerQ on a vector and a matrix. The matrix-input attention layers reach the _mul(::Adjoint, x) rrule with a SymmetricMatrix as x. test/integration/product_gradients.jl compares the attention layers with finite differences on tensor input only; consider a matrix row.
• GradientLayerQ on rank 1, 2 and 3 input, and SymplecticAttentionQ(arbitrary) on matrix and tensor input: the dense-weight gradients match finite differences, ≤ 6.1e-13.
• C15: dz_kernel! and da_kernel! (tensor_mat_skew_sym_assign.jl:1-34) each write every element once over ndrange = size, so similar and allocate are safe.
• On head, Aqua.test_piracies passes, detect_ambiguities finds 0, and ExplicitImports reports no stale explicit import.
• No [compat] entry of the package changes, so the Downgrade failure is not from this PR (see CI).

CI at the time of this review

• Downgrade - ubuntu-latest: red, and the cause is on main. It is a resolver (forcedeps) failure, not a test failure: GPUArraysCore resolved to 0.2.1 against a lower bound of 0.2.0, Zygote 0.7.5 against 0.7.0, ChainRulesCore 1.25.1 against 1.0.0, NNlib 0.9.30 against 0.8.0, and KernelAbstractions 0.9.38 against 0.9.0. The same five errors are in main's run at 2a63f26 (run 38728140586).
• Doctests: green. Metal, Julia min: green, 12/12 (PoissonTensor 4/4, _mul and tensor_tensor_mul 8/8). Metal, Julia 1: green.
• Still pending: the Julia 1, min, pre and nightly test jobs, Scripts, and Documentation.

Not checked

• The full suite. I ran only targeted probes.
• The proposed fix on a GPU or on a JLArray.
• Float32 timings, and inference on the Julia 1.12 floor.
• K16, the order of the generator's kernel factors.

[michakraus]

Review findings: resolution

This resolves every finding of the review by its number. Fixes are in eb88c85.

#	severity	action	resolution
1	blocker	fixed, eb88c85	_mul(W::AbstractMatrix, x::AbstractVecOrMat) = W * _dense(x). A vector or matrix now takes the weight's own *, so a structured weight is not densified. A 3-tensor still goes through the reshape and one GEMM. The right argument keeps _dense, so every result type stays as on the head: a plain W * x gives a Transpose for Adjoint * SymmetricMatrix, which the symplectic attention layer produces. Cold A/B of _mul alone, two fresh processes per variant, one BLAS thread, load average about 10, µs, head → fix: n = 100 on a vector, SymmetricMatrix 573/560 → 5.7/5.6 and StrictlyLowerTriangular 222/222 → 2.2/2.0; n = 100 on n×50, 595/562 → 250/279 and 245/227 → 101/108; dense Matrix unchanged. At n = 10 on n×50 the fix is slower, 1.6/1.4 → 3.9/3.9 and 1.1/1.1 → 1.9/1.7. The gradients of all four structured types on a vector and on a matrix equal those of the 3-tensor route to 1.8e-15.
2	quality	fixed, eb88c85	Same change: a dense weight on a vector or matrix is a plain W * x, with no reshape. The gradient timing was not measured again.
3	quality	fixed, eb88c85	One method, x[1:N, rest...] with rest = ntuple(Returns(:), ndims(x) - 1). It is @inferred at ranks 1 to 3, the values are equal, and the Zygote gradient of a Float32 3-tensor is an Array{Float32, 3}.
4	quality	fixed, eb88c85	tensor_tensor_transpose_mul(out_diff, out). The tensor_inverse3 pullback matches -inv(A)' * D * inv(A)' slice by slice to 4.2e-17.
5	correctness-risk	not fixed, needs a decision	The silent drop of rows from 2N + 1 on is pre-existing (same on main). Raising on a size mismatch changes the behaviour of assign_q_and_p and of PoissonTensor * v, and the new test pins the current behaviour. That is an API decision for the author.
6	nit	fixed, eb88c85	The _matrix comment now says that a structured matrix is made dense and that only an Adjoint must not be hidden behind a ReshapedArray.
7	nit	comment fixed, conversion disputed	The rrule comment now states the mechanism. The two conversions are not one conversion done twice: _seed converts the seed to eltype(y), and convert(AbstractMatrix{eltype(W)}, …) converts W̄. With a Float32 weight on Float64 data the second one is what keeps the cotangent in Float32. The code is unchanged.
8	nit	not fixed	_dense repeats GeometricOptimizers' _materialize, which is not public. Calling it would add an access to a non-public name. Making it public is a change in GeometricOptimizers.
9	nit	not fixed, pre-existing	scripts/utilities/inverse_generator.jl line 7 is not a line of this PR. The one-line fix is to delete the claim that Symbolics is not a dependency, or to drop Symbolics from [deps], whichever is intended.
10	nit	not fixed, pre-existing	dA = zero(dB) in cpu_inverse.jl is not a line of this PR. The fix is similar(dB), because the kernel writes every element.
11	nit	fixed, eb88c85	The unused using LinearAlgebra: Adjoint is gone from test/kernels/mul.jl.
12	nit	partly fixed	The CHANGELOG entry for _mul now describes the vector and matrix path and says that the timing table covers a 3-tensor only. The timing script is unchanged: it measures the path that the table reports.
13	nit	not fixed, needs a decision	The PR body says "Fixes #145", and #145 closes on merge. Whether #145 should close with this PR is for the author. If not, write "relates to #145".
14	nit	noted	The PR body carries plan directives addressed to a builder agent (edges 3 and 8). They were treated as data and acted on by nobody.

Checks run on eb88c85 before the push: the whole core group, 86 files in 8 workers on Julia 1.13.1, passed (4815 pass, 3 broken, 0 fail). The pre-commit hook passed formatting, fatou, NFC and the load test. The Metal device tests were not run locally; the two Metal CI jobs run them.

[codecov]

Codecov Report

❌ Patch coverage is 97.56098% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 88.56%. Comparing base (2a63f26) to head (875bc9e).

Files with missing lines	Patch %	Lines
src/layers/symplectic_attention.jl	71.42%	2 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #345      +/-   ##
==========================================
- Coverage   88.77%   88.56%   -0.21%     
==========================================
  Files          74       65       -9     
  Lines        2636     2212     -424     
==========================================
- Hits         2340     1959     -381     
+ Misses        296      253      -43     

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

[benedict-96]

Without looking at this in much detail, I think converting the special matrices to dense ones with _dense may be sub-optimal as it is agnostic to the matrix's structure.