Keep a Float32 network Float32: training history, _norm, PoissonTensor defaults, the optimizer step size and two parameterlength counts

[michakraus]

Part B of a whole-package audit of GeometricMachineLearning at v0.7.0: a Float32 network must stay Float32. Six places did not. As with Part A, Pkg.test() is green on main with every one of them in place, so each fix lands with a test that fails before it and passes after, verified per test rather than per batch.

The house pattern is already here — src/loss/symplectic_euler_loss.jl and src/loss/variational_midpoint_loss.jl convert the timestep to eltype(input) rather than promoting against it, and src/architectures/hamiltonian_neural_network.jl builds its Poisson tensor from Int for the same reason. This makes the rest of the package follow it.

What is fixed

A Float32 network's training history came back Vector{Float64}. optimize_for_one_epoch! was careful to accumulate in T, and then the Optimizer functor stored the result in an untyped zeros(n_epochs).

_norm widened two of its three methods. The (q, p) arm divided by √2 and the generic NamedTuple arm by √length, both Float64, while the AbstractArray arm was already right. Three arms of one function disagreeing. It reaches users through reduction_error and projection_error, so a Float32 reduced system reported a Float64 error.

PoissonTensor carried an undocumented CPU/GPU element-type split. PoissonTensor(backend, n2) defaulted to Float32 and PoissonTensor(CPU(), n2) to Float64 — added in one 2024 commit with no stated reason. Both are removed rather than unified: Apple GPUs have no Float64 at all (MtlArray(rand(Float64, 4)) errors), so a shared default could only be wrong for one side. Every call site in scripts/, docs/ and test/ that names a backend already names a T, so a caller now has to. PoissonTensor(n2) is unaffected and keeps its Float64 default — which the docstring now actually states.

An untyped zeros in a DataLoader constructor, whose three sibling methods all wrote zeros(T, …).

A Float64 step size reached GeometricOptimizers._rmul! in the GO-native leaf step, where the three _euclidean_update! methods all convert with T(step_size). The funnel stays Float64 deliberately; the conversion belongs at the point of use.

This one is not a style point. rmul!(::MtlArray{Float32}, ::Float64) raises InvalidIRError compiling GPUArrays.gpu_rmul_kernel!, and with a Float32 scalar it returns normally — so GPU training did not run for any architecture. Confirmed on an Apple GPU: on this branch, a GSympNet on MetalBackend() trains two epochs and returns a finite Vector{Float32} with its parameters still MtlMatrix{Float32}; the same call on main raises the error. Reading the rest of _leaf_optim_step! confirms this was the only Float64 scalar reaching a kernel there.

Two parameterlength functions routed integer counts through Float64 division and back through Int(...). Exact at the sizes in use, wrong beyond the Float64 mantissa. Rewritten with integer division: 3M(M + n) is always divisible by 2n because n_heads ∣ M, and M2(M2 + 1) is always even, so the rewrites are algebraically identical wherever the old ones were exact.

Recorded, not fixed

src/layers/classification.jl's average = true and average = false methods disagree on output element type; the decision and its reason are in the changelog. A new open issue B10 records that the DataLoader(::EnsembleSolution) method carrying the zeros fix appears unreachable: every GeometricEquations 0.21.3 equation type rebuilds its own ics inside initialstate, so nothing produces the two-key (:q, :v) shape it dispatches on, and neither GeometricSolution nor EnsembleSolution offers a field-based constructor. The method is left in place — removing dead code is a separate change.

Out of scope

src/arrays/poisson_tensor.jl's Base.:* methods are untouched; they account for 17 of the 23 ambiguities and are a separate change.

Pre-PR verification

Fixed during verification

• _norm(::NamedTuple) no longer throws on a nested NamedTuple. The element-type fix had used promote_type(map(eltype, values(dx))...), which yields an array type for a nested argument, so T(length(dx)) became Vector{Float32}(2). Measured on (q = (a = …,), p = (b = …,)): main returns 1.285065671879467, that version raises MethodError: no method matching Vector{Float32}(::Int64), and typeof(sum(map(norm, dx))) returns 1.2850657. Same element type, same value, generality restored.
• PoissonTensor's docstring now lists all three constructors and states that the element type defaults to Float64 and the backend to CPU().
• Five comments and two testset names rewritten from narrating the replaced code to stating what the test asserts, and one optimizer.jl:310 reference dropped — a line number is exactly what this branch shifts elsewhere.
• Nine references the branch's own line shifts falsified: test/aqua.jl's three piracy witnesses, and in ## Open Issues those three plus add! and the three PoissonTensor ambiguity witnesses. All re-measured.

Checked and clean

• Apple GPU, on hardware: PoissonTensor(MetalBackend(), 4, Float32) works, PoissonTensor(MetalBackend(), 4) is now a MethodError, PoissonTensor(4) is Float64 on the host. No backend can default to a type Metal cannot hold, and no two-argument backend call site remains anywhere.
• Piracy: Aqua.Piracy.hunt → 12, matching the gate; none in a changed line.
• Ambiguities: 23, unchanged — this branch does not touch that set.
• Inference: _norm returns Float32 on all three arms with no Any; both parameterlength methods return Int64; _leaf_optim_step! infers Nothing over 807 statements with no Any or Box; the Optimizer functor now infers Vector{Float32}.
• Allocations, fresh process at --check-bounds=auto: _norm allocates 0 bytes before and after, on both arms. No hunk adds an allocating construct.
• Both large parameterlength values re-derived against BigInt: MHA 19296855159637518 against the Float64 path's 19296855159637520, PSD 11427319538133785 against 11427319538133784.
• jldoctest blocks in the changed files re-run by hand, output unchanged. JuliaFormatter clean on every changed file. NFC clean.
• Full suite after the verification edits: 69 testsets, 4143 passing assertions against 4122 on main, 2 pre-existing broken, 0 failures.

Pre-existing, next to the diff — not fixed

• Two "There used to be…" comments at src/optimizers/optimizer.jl:275 and src/utils.jl:58 narrate history outside the changelog.
• fatou lint: typeof(f) == NothingFunction at src/utils.jl:4 should be an isa test.

Not checked

• The docs/make.jl build. The jldoctest blocks in the changed files were re-run individually instead.

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

[codecov]

Codecov Report

❌ Patch coverage is 90.90909% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 80.93%. Comparing base (28e2537) to head (0f300a2).
⚠️ Report is 19 commits behind head on main.

Files with missing lines	Patch %	Lines
src/data_loader/data_loader.jl	0.00%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #295      +/-   ##
==========================================
+ Coverage   80.61%   80.93%   +0.32%     
==========================================
  Files          80       80              
  Lines        3162     3163       +1     
==========================================
+ Hits         2549     2560      +11     
+ Misses        613      603      -10     

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

Verdict: request changes — two file:line citations that the rebase invalidated are still wrong, and both sit in the part of the branch whose job is to keep citations pointing at the right line. Everything else in the diff verifies: the five element-type fixes, every number the tests quote, the ambiguity and piracy baselines, and all three conflict resolutions.

Reviewed at a07aad2e, which is the current PR head (gh pr view 295 --json headRefOid and git ls-remote origin refs/heads/fix/audit-part-b both give a07aad2e77b8987ae9eb6ac73774066cf8deea03). CI: all eleven jobs pending on that head at review time (run 35358964036), so no CI job was read. The test results below are local runs, named individually.

Findings

#	severity	file:line	claim	evidence
1	bug	CHANGELOG.md:3153, CHANGELOG.md:3179	B10 cites src/data_loader/data_loader.jl:323 for the method and :337 for the fix. Both are off by two after the rebase; the real lines are 325 and 339.	grep -n in the worktree: 325:function DataLoader(ensemble_solution::EnsembleSolution{T, T1, Vector{ST}}; and 339: data = zeros(T, sys_dim, input_time_steps, n_params). Pre-rebase they were 323/337; b7f3b66e on main ("Throw instead of returning nothing on an unexpected autoencoder keyword") added two lines above them.
2	bug	test/aqua.jl:39	Still cites poisson_tensor.jl:39 for getindex(::PoissonTensor, i, j). This branch moved that definition to line 42, and CHANGELOG.md:3126 was re-pointed to :42 — the two documents now state different lines for the same method.	grep -n "Base.getindex" src/arrays/poisson_tensor.jl gives 42:. grep -n "poisson_tensor.jl:" test/aqua.jl CHANGELOG.md gives test/aqua.jl:39: … poisson_tensor.jl:39 against CHANGELOG.md:3126: … poisson_tensor.jl:42.
3	quality	CHANGELOG.md:655	Cites src/data_loader/data_loader.jl:367 for the EnsembleSolution constructor a script uses; it is at 369. Pre-existing — already stale on origin/main, from the same two-line shift as finding 1. It belongs in this branch because this branch is the one re-pointing citations into that file.	git show origin/main:src/data_loader/data_loader.jl | grep -n "^function DataLoader(ensemble_solution" gives 325: and 369:.
4	bug	src/layers/grassmann_layer.jl:28	N > M ? (N - M)*M : (M - N):N — a colon where a * belongs. parameterlength(::GrassmannLayer) returns a UnitRange, not an integer, whenever N <= M. Pre-existing, but it is the third parameterlength of the set this branch audited and the only one that is outright broken.	In a live session: parameterlength(GrassmannLayer(10, 4)) returns 6:5, of type UnitRange{Int64}. The sibling at src/layers/stiefel_layer.jl:20 already uses ÷ and is correct.
5	correctness-risk	src/data_loader/optimize.jl:95	zeros(eltype(dl), n_epochs) throws InexactError on the first write when eltype(dl) <: Integer, where zeros(n_epochs) did not. An integer DataLoader is constructible.	eltype(DataLoader(rand(1:10, 4, 16); autoencoder = true)) == Int64. I could not reach the failure: the same training call dies earlier with MethodError: no method matching (::FeedForwardLoss)(::Chain{…}, …), because the loss rejects integer input. So no observed regression — worth a line in the entry rather than a code change.
6	nit	test/layers/psd_parameterlength.jl:18	The "ordinary sizes" testset recomputes expected with the same ÷ expression as src, so it cannot fail on an arithmetic error — only on an edit to the formula. MultiHeadAttention does have an independent cross-check; PSDLayer has none.	test/architectures/check_parameterlengths.jl:22 derives the MHA count independently as 3*N÷n*Int(n*(N-(n+1)/2)), and passes under the rewrite (run below). No equivalent exists for PSDLayer.
7	nit	CHANGELOG.md, the _norm bullet	The entry says the fix "reaches users through reduction_error and projection_error", but test/reduced_order_modeling/norm_eltype.jl only asserts _norm itself. The user-facing claim has no test.	src/reduced_system/reduced_system.jl:217,218,237,238 are the four _norm call sites; no test asserts the element type of either error function.

The three conflict resolutions — all correct

1. test/data_loader/runtests.jl — both @safetestset blocks are present: main's autoencoder_keyword.jl and this branch's training_history_eltype.jl. Both files exist, both are reached, both pass. test/reachability.jl, which walks the include closure of runtests.jl, passes with an empty allowlist.
2. CHANGELOG.md, "Removed (breaking)" — main's legacy/hnn/ + legacy/mtk/ entry and the branch's PoissonTensor entry are both present, neither truncated. git diff origin/main...HEAD -- CHANGELOG.md | grep '^-' returns exactly six lines, every one of them an old citation replaced by the re-pointing. No prose was dropped.
3. CHANGELOG.md, B8 — main's "eighteen … of the 23 this entry originally reported" text is kept and is internally consistent: 23 − 5 triaged = 18, and 17 * pairs + 1 Dense/Affine = 18. The re-pointed citations are right. src/arrays/poisson_tensor.jl:75,78,81 are the three Base.:*(::PoissonTensor{T}, …) methods (main had them at 71/74/77; this branch adds four lines above them), and poisson_tensor.jl:42 is Base.getindex. Both checked with grep -n against HEAD and against origin/main.

Verified and fine

• Every other file:line the branch adds or shifts resolves. src/utils.jl:53 is add!, :65 and :71 the two Base.:+ pirate methods, :167 the Base.:≈ method. Main had 49/61/67/163; the _norm rewrite adds four lines above them. Checked in both CHANGELOG.md and test/aqua.jl. src/utils.jl:19-27 (CHANGELOG.md:596) and src/optimizers/optimizer.jl:21 (CHANGELOG.md:3109) are unshifted and still correct.
• B10 is free. No **B10. anywhere in CHANGELOG.md on origin/main, and a sweep of every refs/remotes/origin/* finds it only on this branch. The highest number on main is B9. The collision risk is forward-looking only: if the PositionalEncoding follow-up allocates B10 independently, two different defects share one number. Reserve it now, or have the follow-up start at B11.
• The baselines test/aqua.jl asserts still hold, measured in a clean process. julia --project=test -e 'using GeometricMachineLearning, Aqua, Test; …' gives detect_ambiguities(GML) = 18 and piracies = 12. That matches the @test length(Aqua.Piracy.hunt(…)) == 12 gate and the "18 when GML is loaded alone" comment. Removing the two PoissonTensor constructors moved neither number.
• Every number quoted in the new tests reproduces. parameterlength(PSDLayer{302355014, 302355018}) gives 11427319538133785; BigInt arithmetic gives the same; the old Int(M2 * (N2 - (M2 + 1) / 2)) path gives 11427319538133784 — exactly as the comment says. parameterlength of the M = 83767764, n = 6 MultiHeadAttention gives 19296855159637518, BigInt agrees, and the old path gives 19296855159637520. For the step-size guard, replaying _leaf_optim_step! with the pre-fix _rmul!(direction(cache), step_size) line separates the two arms by 2.9802322f-8 for both Adam() and MomentumMethod(0.5) at seed = 2, step_size = 0.1 — the figure the file's header claims. All four new tests are real regression guards.
• The element-type fixes do what they say. Cold process, --check-bounds=auto: _norm((q, p)) and _norm((a, b)) on Float32 arguments return Float64 under the old expressions and Float32 under the new ones, with 0 allocations both before and after. eltype(PoissonTensor(4)) == Float64, and PoissonTensor(CPU(), 4) is a MethodError.
• T(step_size) matches the file's own pattern. The three _euclidean_update! methods at src/optimizers/optimizer.jl:259,267,282 all convert first, and the funnel really does always hand a Float64: _step_size(η::Real, ::Int) = Float64(η) at :204, the DecayingStatic arm at :209, _optimizer_step_size(η::Real) = Float64(η) at :213, and _step_size_from_linesearch(ls::Static) = Float64(ls.α) at :247.
• The PoissonTensor removal breaks no call site. grep -rn "PoissonTensor(" src docs scripts test: every call that names a backend already names a T (test/arrays/poisson_tensor.jl:17, and the docstring's own jldoctest at src/arrays/poisson_tensor.jl:26). The entry's claim holds verbatim.
• The cross-package claims check out. AbstractNeuralNetworks 0.8's _norm(dx::NamedTuple) = sum(apply(norm, dx)) / √length(dx) (losses.jl:69 in the resolved copy) does carry the identical defect. GeometricEquations is at 0.21.3 in both manifests, and initialstate(equ::SODE, t, ics, params) at src/odes/sode.jl:152 does return (q = _statevariable(ics.q, periodicity(equ)),) alone, while IODE at src/odes/iode.jl:188 returns (q, p, v). ClassificationLayer's two doctests do show Matrix{Float64} at src/layers/classification.jl:43 and Matrix{Int64} at :60.
• Type stability and allocations. Base.return_types gives Float32 for both NamedTuple arms of _norm and Int64 for both rewritten parameterlength methods. Base.eltype(::DataLoader{T}) where {T} = T at src/data_loader/data_loader.jl:511, so zeros(eltype(dl), n_epochs) is inferable. No allocation change is attributable to the diff.
• Tests run, and what they returned. julia --project=test over reachability.jl plus the arrays, attention, layers, optimizers, reduced_order_modeling and data_loader drivers: exit 0, no failures, no errors, none broken. Notable summaries — Reachability of every test file 2/2, Symplectic Potential (array tests) 43/43 (which includes the two new PoissonTensor assertions), parameterlength(::MultiHeadAttention{M,M,true}) is exact 4/4, parameterlength(::PSDLayer) is exact 5/5, The GO-native leaf step scales in the parameter's own element type 2/2, _norm keeps the element type of its argument 6/6, A Float32 network's training history stays Float32 4/4. Separately, Check parameterlength 9/9 — the independent cross-check on the MHA rewrite.
• Hygiene. All nineteen changed .jl files pass JuliaFormatter.format(f; overwrite = false) under the repository's sciml config. Every changed file, CHANGELOG.md included, satisfies s == Unicode.normalize(s, :NFC). No new exported name and no new public API, so nothing to register in the docs; the PoissonTensor docstring gains three signatures and leaves the jldoctest body untouched.
• Comment tense. The new src comment at src/arrays/poisson_tensor.jl:60-63 states an absence in the present tense and gives a design reason (Metal has no Float64) rather than narrating a removal. The test headers describe counterfactuals, not history, and each counterfactual was reproduced above.

Not checked

• CI. Every job was still pending. Nothing above reads a CI job, in either direction.
• The docs build and the doctests. Pkg.test() does not run jldoctest blocks, and I did not run docs/make.jl. The PoissonTensor doctest body is unchanged, so the risk is low, but it is unmeasured. The Doctests and Documentation jobs are the check.
• The full Pkg.test(), and the losses, transformers, kernels, activations, parameters and docstrings drivers. I ran the six drivers the diff touches, plus architectures/check_parameterlengths.jl.
• Windows, macOS and the 1.10 floor. One local run, one platform. Nothing in the diff looks version-sensitive — ÷, typeof(n)(...) and zeros(T, n) are all long-standing constructs — but that is reading, not measuring.
• B10's segfault paragraph. I did not attempt the two ccall bypasses it reports, so "both crashed the Julia process" is taken on the entry's word. The reachability argument in front of it — that no GeometricEquations 0.21.3 equation type yields a two-key (:q, :v) state — I did check directly for SODE and IODE.
• ExplicitImports and a full Aqua.test_all. The diff adds no import or using to src/, so neither is diff-attributable. I ran only the two checks that have asserted baselines.

What to change before merge

Findings 1 and 2 are one-line edits, and they are the whole of "request changes": point B10 at data_loader.jl:325 and :339, and bring test/aqua.jl:39 to poisson_tensor.jl:42 so it agrees with CHANGELOG.md:3126. Finding 3 is the same class and one more line. Finding 4 — grassmann_layer.jl:28 returning a UnitRange — is a real bug in the third member of the set this branch audited, and reads as belonging here rather than in a follow-up nobody opens.