Make PositionalEncoding trainable, and stop the keyword changing which inputs a Transformer takes

[michakraus]

Answers the review of #294, which merged before these fixes existed. Every finding was reproduced
against the branch first, then fixed, then re-measured.

The blocker: a network with the keyword could not be trained

positional_encoding builds its matrix with a setindex! loop. Zygote traces into it, so every
gradient through a network carrying the layer raised
Mutating arrays is not supported -- called setindex!(Matrix{Float32}, ...). Measured both ways,
for the parameters and for the input:

pe=false gradient wrt params OK      pe=true gradient wrt params ERROR: Mutating arrays ...
pe=false gradient wrt input  OK      pe=true gradient wrt input  ERROR: Mutating arrays ...

The builder takes a type and two lengths, and none of them is a differentiable quantity, so it is
declared ChainRulesCore.@non_differentiable. All four gradients are taken now.

No test in the suite took a gradient through the layer, which is why a green matrix said
nothing about it. The new testset takes both, with the keyword and without.

The keyword added an input type

A vector has no second axis to read a sequence length from, and broadcasting one against the
dim × 1 encoding returned a dim × 1 matrix — which MultiHeadAttention then accepted, although
it rejects the vector itself. So positional_encoding = true accepted an input that
positional_encoding = false refuses, at a different rank:

pe=false Vector -> MethodError        pe=true Vector -> (4, 1)

The functor now takes a matrix or a batch of matrices. A vector raises MethodError either way,
and the test asserts that through the architecture as well as on the layer.

The CPU-only limitation is restated, and recorded

On a device array the broadcast does not slow down — it fails to compile, because a host array
cannot be read from a kernel. Both docstrings say so now; the Transformer keyword carried no
warning at all. The defect is recorded as B11 under Open Issues, because no test can catch it
while the suite has no GPU job. Closing it means giving the builder the backend, and that cannot be
verified without a GPU job to run it under.

B11 rather than B10: #295 already carries a B10, for the EnsembleSolution DataLoader.

Four items of prose

• The CHANGELOG still attributed the zero base of pos to the paper, which gives none — the one
  claim Add a positional encoding layer, and end legacy/ #294's own body listed as false and its docstring already corrected, so the two disagreed
  inside a single release.
• A test comment pointed at "the comment at the top of this file", which is not there.
• Another narrated the branch's history rather than the code's state.
• The allocation figure quoted in Add a positional encoding layer, and end legacy/ #294 was the helper's, not the layer's:
  @allocated positional_encoding(Float32, 32, 64) is 8272 B, while the layer call is
  16544 B, because the sum allocates too.

Checks

The positional-encoding testset passes with 115 assertions, up from 106. Type stability
unchanged: the two functor methods return Matrix{Float32} and Array{Float32, 3}, both concrete.
JuliaFormatter clean on all three changed .jl files; the pre-commit hook's format, lint, NFC and
load stages all passed.

The full Pkg.test() has not been run on this branch. The new @non_differentiable adds an
rrule method, and test/aqua.jl gates on exact piracy and ambiguity counts, so the matrix is the
check that matters here.

🤖 Generated with Claude Code

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[michakraus]

Review

Approve, once the matrix is green. No blocker and no bug. Both central claims reproduce against
the branch, and both were measured on main first. Four nits and two notes follow.

CI is not concluded at the time of writing: no job in the run has a conclusion yet. Since the
body states the full Pkg.test() was never run on this branch, the matrix is the check that
matters, and this approval is conditional on it.

Verified

The AD blocker is real, and it is fixed. Two cold processes, the same script, on main and on
the branch tip:

main    pe=true grad wrt params  ERROR: Mutating arrays is not supported -- called setindex!(Matrix{Float32}, ...)
main    pe=true grad wrt input   ERROR: Mutating arrays is not supported -- called setindex!(Matrix{Float32}, ...)
main    pe=false both            OK
branch  pe=false/true, both      OK   (input gradient size (4, 5), all finite)

The gradient is also correct, not merely non-nothing. Against central_fdm(5, 1) on the pe=true
network the maximum absolute difference is 0.0165, with ‖g‖ = 5.763 against ‖g_fd‖ = 5.755 —
0.3 %, which is Float32 noise — and all(iszero, g) is false.

The input-type widening is real, and it is fixed. On main a vector at pe=true returned a
4 × 1 matrix at the layer and through the architecture, while pe=false raised. On the branch
both raise MethodError, at the layer and through the architecture. The matrix, the (4, 5, 3)
3-tensor and the (q, p) NamedTuple paths all still work.

@non_differentiable is the right instrument here. The arguments are a type and two integers,
and the output is a constant added to the input, so the true Jacobian is the identity — which the
finite-difference check above confirms. The generated rrule and frule are keyed on a
package-owned singleton type, so they are not piracy: Aqua.Piracy.hunt returns 12, the count
test/aqua.jl:63 asserts, and detect_ambiguities(GeometricMachineLearning; recursive = true)
returns 18, the documented baseline, with no entry naming positional_encoding. With Zygote
loaded the new rrule is ambiguous against 0 of 1190 loaded rrule methods, and the new frule
against 0 of 904. test/aqua.jl standalone: 10 of 10 pass.

The transformers group passes locally, with Sinusoidal positional encoding | 115 115 — the body's
figure, and 106 plus the 9 new assertions. Return types are Matrix{Float32} and
Array{Float32, 3}, both concrete. JuliaFormatter is clean on all three .jl files, all four
changed files satisfy s == Unicode.normalize(s, :NFC), and no docstring is detached. The amended
CHANGELOG text sits under ## [Unreleased] — 0.8.0, so it corrects an unreleased entry rather than
rewriting history. No file:line citation anywhere in the repository is falsified by the diff.

All four prose items in the body are genuinely fixed, and both allocation figures reproduce exactly:
@allocated positional_encoding(Float32, 32, 64) is 8272 B, and the layer call is 16544 B.

Findings

#	severity	file:line	finding
1	nit	test/transformers/positional_encoding.jl:126-127, :134	The new comments narrate the fixed defect in the past tense — "every gradient … raised Mutating arrays is not supported", and "the failure was in the forward trace". This is the rule the PR itself applies two hunks earlier, where it deletes "It did:" from the (q, p) comment. The counterfactual at src/layers/positional_encoding.jl:53-55 gets it right.
2	nit	test/transformers/positional_encoding.jl:135	@test Zygote.gradient(...)[1] !== nothing asserts only that an object came back. A rule that returned zeros, or NoTangent, would pass it. The gradient is in fact correct — see above — but the assertion does not pin that.
3	nit	test/transformers/positional_encoding.jl:110-121	The input-type testset asserts a bare Vector only. The review of #294 also measured a (q, p) pair of vectors accepted at pe=true and rejected at pe=false. That case is fixed on this branch — MethodError at the layer and at both keyword settings through the architecture — but nothing asserts it, so it can regress silently.
4	nit	CHANGELOG.md:1030 against src/layers/positional_encoding.jl:66	The CHANGELOG now reads "one dim × seq_length allocation per call, and a second for the sum", while the docstring still reads "costs one dim × seq_length allocation". Both are true — total against marginal — but read side by side they contradict each other, which is the same defect this PR fixes for the zero base of pos.
5	note	CHANGELOG.md:3109	B11 is correct only while #295 is open and merges. If #295 closes, the list jumps B9 to B11 with no B10. The two branches also conflict textually: git merge-tree reports a content conflict in CHANGELOG.md, so whichever merges second needs a rebase.
6	note	—	The review of #294 asked for the host allocation to be fixed, not recorded. This PR records it as B11 and adds the Transformer warning, which is the fallback that review offered. That is a scope decision rather than a defect, and it is left as it stands.

Not checked

The GPU claim in B11 — that the broadcast fails to compile rather than running slowly — is not
verified here, for want of a GPU run. It matches the InvalidIRError quoted in the review of #294.
The local runs covered the transformers group and test/aqua.jl only, under --check-bounds=auto
rather than the =yes that Pkg.test uses below 1.13. The Documenter build and the doctests are
left to CI; the diff changes no jldoctest block.

🤖 Generated with Claude Code

[michakraus]

The head has moved to 669ddcc4. My review above cites 63ac9770, which the rebase replaced.

Since #295 merged, this branch is rebased onto 25282c5e. The CHANGELOG.md conflict was the one
named as finding 5: both branches added a different B10 at the same point in Open Issues.
#295's B10 stays exactly as merged, and the CPU-only defect follows it as B11. The net diff in
that region is additions only, and Open Issues now runs B9, B10, B11.

The branch is two commits rather than three. 04be7140 renumbered B10 to B11 in one line, and the
conflict resolution already writes B11, so that commit had nothing left to do and was dropped. The
net diff is unchanged by this.

Findings 1 to 4 were fixed in the previous head and carry over untouched. Finding 6 is still open.

JuliaFormatter and NFC pass on all four changed files at the new tip. The suite was last run in
full at 63ac9770 — 120 assertions in the positional-encoding testset, Aqua 10 of 10 — and only
CHANGELOG.md has changed since. The matrix on this tip is the remaining check.

🤖 Generated with Claude Code

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 80.94%. Comparing base (0f300a2) to head (669ddcc).
⚠️ Report is 1 commits behind head on main.

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #296   +/-   ##
=======================================
  Coverage   80.93%   80.94%           
=======================================
  Files          80       80           
  Lines        3163     3164    +1     
=======================================
+ Hits         2560     2561    +1     
  Misses        603      603           

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

Review of 669ddcc4 — approve

Approve. Second review, against the new tip. The rebase is clean, the four findings of the
review of 63ac9770 are all closed at this tip, and the matrix that review made its approval
conditional on has now run. No blocker, no bug. Two nits and one note follow, none of them gating.

CI — read from the jobs

Twelve of thirteen checks pass. One fails, and it is the documented advisory failure:

job	result
Julia 1 / min — ubuntu, macOS, windows (6 jobs)	pass
Julia pre — ubuntu	pass
Doctests, Documentation (1h1m), Scripts	pass
codecov/patch, codecov/project	pass
Julia nightly — ubuntu	failure

The nightly log names one failing assertion, and it is B9, not this diff:

Unbound type parameters detected:
[1] var"#iterate#57"(n_points, ::typeof(iterate), nn::NeuralNetwork{<:NeuralNetworkIntegrator}, ics::BT)
Unbound type parameters: Test Failed at .../Aqua/src/unbound_args.jl:38
  Expression: isempty(unbounds)
ERROR: LoadError: Some tests did not pass: 9 passed, 1 failed, 0 errored, 0 broken.

That is exactly the method CHANGELOG.md records under Open Issues as B9, at
src/architectures/neural_network_integrator.jl:98. .github/workflows/CI.yml:48 gives the job
continue-on-error: ${{ matrix.experimental }}, and nightly carries experimental: true. Nothing in
the diff touches that method. The failure is pre-existing and advisory.

Verified against this tip

The full suite was run, which the body says it had not been. Pkg.test() on Julia 1.13.0,
--check-bounds=auto, in a clean detached checkout of 669ddcc4: passed, with
Sinusoidal positional encoding | 120 120. The body's figure is addressed under finding 1.

Both defects reproduce on origin/main (25282c5e) and are gone at the tip. One script, two
cold processes:

main    pe=false params OK          main    pe=true params ERROR: Mutating arrays is not supported -- called setindex!(Matrix{Float32}, ...)
main    pe=false input  OK          main    pe=true input  ERROR: Mutating arrays is not supported -- called setindex!(Matrix{Float32}, ...)
main    layer on Vector        -> (4, 1)            main   nn(pe=true)  on Vector -> (4, 1)
main    layer on (q,p) vectors -> (q=(2,1), p=(2,1))  main nn(pe=false) on Vector -> MethodError

tip     pe=false/true, params and input      all four OK
tip     layer / nn(pe=false) / nn(pe=true), Vector and (q,p) vectors   MethodError, all six
tip     matrix, 3-tensor, (q,p) matrices, (q,p) 3-tensors              all still work

The gradient is right, not merely non-erroring. Zygote.gradient(y -> sum(l(y, NamedTuple())), x)[1]
is exactly ones(Float32, 4, 5). Through a full Transformer(4, 2, 1; positional_encoding = true)
the input gradient agrees with a central difference at h = 1f-3 to a maximum relative error of
3.0e-4, which is Float32 noise.

@non_differentiable costs nothing in Aqua's terms. test/aqua.jl standalone: 10 of 10 pass,
including the exact gate @test length(Aqua.Piracy.hunt(GeometricMachineLearning)) == 12 at
test/aqua.jl:63. The generated rrule/frule are keyed on a package-owned singleton, so they are
not piracy and they add no ambiguity.

Both allocation figures reproduce exactly. Cold process, warmed, --check-bounds=auto:
@allocated positional_encoding(Float32, 32, 64) is 8272 B, the layer call is 16544 B.

Type stability holds. Base.return_types gives Matrix{Float32}, Array{Float32, 3} and
@NamedTuple{q::Matrix{Float32}, p::Matrix{Float32}} for the three functor methods, and @inferred
passes on all three.

The narrowed signature follows an existing convention.
Union{AbstractMatrix, AbstractArray{<:Any, 3}} at src/layers/positional_encoding.jl:118 is the
same shape already used at src/layers/bias_layer.jl:33.

The comment at test/transformers/positional_encoding.jl:142-143 is accurate. With the keyword
on, keys(nn.params) and keys(g) are both (:L1, :L2, :L3), nn.params.L1 is NamedTuple() and
g.L1 is nothing.

Prose and hygiene. All comments added by the diff are in the present tense, including the two the
earlier review flagged. JuliaFormatter is clean on all three .jl files, and all four changed files
satisfy s == Unicode.normalize(s, :NFC). Both CHANGELOG hunks sit under ## [Unreleased] — 0.8.0
and in the Open Issues appendix, so nothing released is rewritten; Open Issues now reads B9, B10,
B11 with no gap and no duplicate. The diff adds no file:line citation and falsifies none. Zygote
enters the test file as import, not a bare using, and is declared in test/Project.toml with
Zygote = "0.7".

Findings

#	severity	file:line	claim	evidence
1	nit	PR description	The body still says the testset passes with 115 assertions. It is 120 at this tip.	Full Pkg.test() at 669ddcc4: Sinusoidal positional encoding | 120 120. The same file on origin/main gives 106. The follow-up comment already states 120, so only the body is stale.
2	nit	test/transformers/positional_encoding.jl:153-159	The four assertions do not pin the gradient through the encoding. A rule that returned a zero input cotangent passes all of them.	nonzero_somewhere(gradient_wrt_params) is satisfied by L2/L3 alone, because the encoding is the first layer and its pullback does not reach the parameters behind it. size(...) == size(x) and all(isfinite, ...) are both true of zeros(Float32, 4, 5). One line pins it: @test Zygote.gradient(y -> sum(PositionalEncoding(4)(y, NamedTuple())), x)[1] == ones(Float32, 4, 5) — measured true here.
3	note	CHANGELOG.md:3213	B11 records the CPU-only defect rather than fixing it. This was finding 6 of the previous review and is unchanged.	A scope decision, not a defect, and the fallback the review of #294 offered. Closing it needs the backend in the builder and a GPU job to verify under, neither of which this PR can supply.

Not checked

• The GPU claim in B11 — that the broadcast fails to compile rather than running slowly. No GPU
  run is available here. It is consistent with the InvalidIRError quoted in the review of #294.
• ExplicitImports over src/. The diff adds no using or import to src/, so any finding
  would be pre-existing rather than diff-attributable.
• --check-bounds=yes. The local run used auto, which on Julia 1.13 is what Pkg.test()
  inherits. The three min jobs on CI cover yes, and all three pass.

🤖 Generated with Claude Code