Repository navigation
GML part M4: one row-pivoted Cayley kernel at every sequence length, with the closed forms for N = 2 to 5 kept - #347
Merged
Merged
Conversation
…re decided `tensor_cayley` in `src/kernels/cayley.jl` solves `(I + C)Λ = (I - C)` slice by slice by row-pivoted Gauss–Jordan elimination, on one work item per slice, with the loop bound read from `size(C, 1)` at run time. Its rrule reads `C̄ = -(I + Λ)ᵀ Λ̄ (I + Λ)ᵀ / 2` from `Λ` alone: no solve, no identity tensor, and no `Zygote.pullback` inside a kernel. The docstring carries the contract, including the orthogonality bound the tests measure. The layer's generic activation runs the kernel, so every sequence length runs on a device; the closed forms and their dispatch by sequence length stay for now, because the benchmark decides which of them go. `scripts/utilities/cayley_timings.jl` is that benchmark: one cold process per variant, per `N` in 2 to 8 and per element type, BLAS on one thread, a warm-up call and the median of five, forward and backward. The tests are `test/kernels/cayley.jl` (the skew premise of both layer paths, the ratio and absolute orthogonality clauses on three input sets at two scales, the value against a `BigFloat` solve, and the pullback against a central difference), `test/layers/volume_preserving_attention.jl` (the forward pass against an independent host computation and the gradient, on the host and on a `JLArray`, at N = 2 to 8 and with `SL = 0`), and rows in `test/devices/jlarray.jl`, `test/devices/metal.jl`, `test/quality/jet.jl` and `test/kernels/kernel_pullbacks.jl`.
…ver the batch inverse The benchmark of check 1 ran at 76cee62, the commit before this one, which held the kernel and every closed form. Its medians, two runs, are in `CHANGELOG.md` with the machine and the load. Read against the rule of the part's section: - every closed form for N = 2 and N = 3 stays, because a forward median is above 1 for each of them: N = 2 at Float64 forward 1.46 and 1.33, N = 3 at Float32 forward 1.71 and 1.79 and Float64 forward 1.23 and 1.68; - N = 4 and N = 5 have a forward median above 1 as well (1.13 and 1.27 in the first run, 1.13 and 1.23 in the second at Float64, and 1.70 and 1.42 at Float32 in the second), which is the case the section leaves to the user: the closed forms stay until that ruling, and the question is in the pull-request body; - N = 6 to 8 are replaced whatever their ratio, which was the whole point of the kernel: the batch inverse called `inv` and `Zygote.pullback(inv, ·)` on one slice inside a kernel. So `cpu_inverse`, its rrule and `cpu_tensor_cayley` go, and nothing calls them: the layer's generic activation runs `tensor_cayley`. The identity tensor moves from `src/kernels/exponentials/` to `src/kernels/identity_tensor.jl`, which is what the file holds, and C18 leaves `KNOWN_ISSUES.md`. `scripts/utilities/inverse_generator.jl` names `tensor_cayley` as the fallback. The tests that covered the deleted functions lose their rows, in `test/kernels/tensor_inverse.jl`, `test/kernels/inverses/test_cayley_transforms.jl`, `test/kernels/kernel_pullbacks.jl` and `test/quality/jet.jl`, and the rows for the sizes that stay are untouched.
`B` and `dC` were the buffers of the `cpu_inverse` JET rows, and nothing reads them now. `fatou lint` reports the two unused bindings, and the tree keeps `test/` free of its findings.
…tal column The Metal column is measured after all: `Metal.functional()` is true in this machine's bash processes through the test environment, its device is an `AGXG16CDevice` on an Apple M4 Max, and `test/devices/metal.jl` passes 40 of 40 there — the 12 assertions of the two testsets it had and the 28 of the new one. The N = 2 backward ratio of 5.56 on Metal is the row check 1's rule names.
The user kept tensor_cayley4 and tensor_cayley5, so the CHANGELOG entry states the reason for every closed form that stays, with the check 1 medians each rests on, and no longer says that N = 4 and 5 await a ruling. The header of cayley_timings.jl says why Metal has no reference for N = 6 to 8, in place of the claim that Metal has no closed form to compare against.
… benchmark times a device Round 1's three critics fail the part on three things, and this commit answers all of them. The exact skewness of the `:arbitrary` path is now a property of one function the layer calls, `arbitrary_operand`, rather than of a copy of its two lines in `test/kernels/cayley.jl`. A mutant that rebuilds that path to be skew only to rounding now fails the premise row — the two mutants the critics wrote for it are both CAUGHT — where before it survived, because the row asserted on the test's own copy and never ran the layer's code. The row runs at M = 2, 3 and 5, the sizes whose `√M` is not exact, so the mutation's second rounding shows. `scripts/utilities/cayley_timings.jl` takes two optional arguments: `array`, the constructor that moves the batch and the seed onto the device, and `sync`, the wrapper that makes a timed call a time rather than a launch. The header's Metal call is a function, `x -> (Metal.synchronize(); x)`, because `Metal.@sync` is a macro and cannot stand as a value, and the call runs. The Metal column of the `CHANGELOG.md` entry is that script's own `measure` with `MtlArray`, measured twice on an `AGXG16CDevice`, so edge 13(b) holds for it too; the file runs its table only as the program, so including it exposes `measure` and runs no driver. Check 3 measures `‖ΛₖᵀΛₖ − I‖∞`, the matrix infinity norm of the slice that the check names, where it measured the largest entry, and its value rows cover the `:skew_sym` input set as well. Check 4's strict rows run at every `N` in 2 to 8 on `Array` and on a `JLArray`. Check 5's fixed-length cases are `VolumePreservingAttention(4, N)`, and the default layer's gradients join them, on the host and on the device. The gradient rows compare against the value row's own bound, without the factor 10 that nothing explained. The pivot swap moves both members through one loop over the pair. `KNOWN_ISSUES.md` gains K19: far outside check 3's scales, `Float32` returns non-finite slices for a skew `C` of norm 1e8, where the closed form that stays is non-finite on 197 of the same 200 slices and the dense solve's own orthogonality error is 7.3e4, so no method is accurate there.
…at compare against it The check 3 bound `10 n eps(T) (1 + ‖Cₖ‖₂²)` was written out in four places: `bound_per_slice` of `test/kernels/cayley.jl`, the `JLArray` rows, twice in the Metal rows, and the layer test's `tolerance`. `test/helpers/cayley_bound.jl` now holds the one copy, as `test/helpers/eltypes.jl` holds the one list of precisions, and those four files include it. The bound is part of the kernel's docstring, so four spellings of it drift apart from that docstring and from each other.
…ures `test/kernels/cayley.jl`: `scale_to` converts the scale to `T` before it multiplies a set, so a row of a precision draws inputs of that precision. The `Float32` `:arbitrary` rows measured a `Float64` input before, because `1e2` is a `Float64`, and each input set now asserts its element type, so a promotion cannot return unnoticed. `src/kernels/cayley.jl`: the cotangent's addition, negation and division are one broadcast after the batched product, so the last step of the `rrule` allocates one array and not three. The docstring states the sizes and scales the tests cover, the ratio's scope and the matrix whose symmetric part is positive definite: the bounds are measured bounds, and elimination can overflow where the tests do not reach. `src/layers/volume_preserving_attention.jl`: the docstring names the generic activation, so it does not read as if the four closed forms that stay had gone. `scripts/utilities/cayley_timings.jl`: the untimed warm-up call is synchronized as well as the timed ones, so no unfinished device work can be counted into the first sample. The header describes what the script measures rather than where the table beside it came from. `KNOWN_ISSUES.md`: K19's reproducer carries its imports and its seed, so the lines around it run as they stand. `CHANGELOG.md`: check 3's worst ratios, taken again on the corrected rows, are 1.91 for the hand-built skew tensor, 1.57 for the `:arbitrary` path and 1.34 for the `:skew_sym` path; the worst exactly skew fraction of the bound is 0.11, and the `:skew_sym` reference's excess over that bound is 3.58. The entry stated figures from the promoting path. Its headline names the kernel as the generic path, because the closed forms for N = 2 to 5 keep their own dispatch.
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #347 +/- ##
==========================================
- Coverage 91.29% 86.72% -4.58%
==========================================
Files 65 58 -7
Lines 2080 1341 -739
==========================================
- Hits 1899 1163 -736
+ Misses 181 178 -3 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
michakraus
commented
Oct 17, 2026
michakraus
left a comment
michakraus
left a comment
Member
Author
There was a problem hiding this comment.
Review of ef1fbae62a5b0739e783a7dfef25e26dc6f38061
Verdict: request changes. (This is posted as a comment review, because the author's own account cannot request changes.)
The kernel and its rrule are numerically correct. But the benchmark that decides which closed forms stay compares unequal code. The kernel has no @inbounds, and the closed forms it is timed against do. With @inbounds, the kernel's CPU forward pass is 1.5–2.3× faster. At that speed the PR's own rule deletes at least tensor_cayley2. The new code can also be made considerably shorter.
Findings
| # | severity | file:line | claim | evidence |
|---|---|---|---|---|
| 1 | blocker | src/kernels/cayley.jl:55-104 |
The kernel body has no @inbounds. Every closed form it is benchmarked against has one (src/kernels/inverses/inverse_2x2.jl:3, inverse_3x3.jl:3, inverse_4x4.jl:3, inverse_5x5.jl:3). So check 1, and the reasons in CHANGELOG.md:154-156 for keeping 841 lines of generated inverses, rest on an unequal comparison. The indexing is in bounds by construction: i, j ≤ size(C, 1) = size(C, 2), and KernelAbstractions masks k to ndrange. |
I added @inbounds begin … end around the unchanged shipped body and changed nothing else. I ran a fresh process for each variant, N and T, on Julia 1.13.2 and KA 0.9.45, with BLAS on one thread, 1000 slices, the median of 41 calls, and a load average of about 5. Kernel forward, shipped → @inbounds (µs): F32 49.4→21.3, 78.7→42.0, 127.2→75.5, 196.9→127.3, and 546→334 at N=8. F64 52.8→24.4, 85.3→46.8, 134.8→86.8, 203.3→135.5, and 563→360. Ratio of the @inbounds kernel to the closed form, forward, over two cold runs: N=2 0.78/0.81 (F32) and 0.79/0.76 (F64). N=3 1.01/0.82 and 0.80/0.78. N=4 0.94/0.94 and 1.04/1.07. N=5 0.85/1.03 and 0.81/1.03. Backward: 0.48–0.56 at every N. On Metal (AGXG16CDevice, warm session, Float32, median of 31 synchronized calls, two runs), shipped → @inbounds forward was N=4 731→551 and 597→441, and N=5 1441→981 and 1498→840. The ratio to the closed form was N=2 0.48/0.65, N=3 0.50/0.59, N=4 0.80/1.01, and N=5 1.33/1.35. Under edge 6's rule, N=2 then has no median above 1 on any column, and N=3 has one median at 1.01. Re-run check 1 with the fixed kernel before you decide the outcome. |
| 2 | bug | KNOWN_ISSUES.md:250-254 |
K19 gives the wrong mechanism. It says that the elimination's intermediates leave the Float32 range, "rather than a small pivot", and that in Float64 they reach 4.6e183 and 2.5e207, against a smallest pivot of 1.2e8 and 2.6e8. The opposite is true. All 23 non-finite slices have an exactly zero last Float32 pivot. The same elimination in Float64 stays below 7e8, and its last pivot is 1.9–8.9. At ‖C‖ ≈ 6.5e8, eps(Float32)·‖C‖ ≈ 78 exceeds the true last pivot. For odd N, C alone is singular, so that pivot cancels to 0. |
The K19 reproducer (Random.seed!(3), N = 5, 200 slices), with the kernel's algorithm traced per slice. Slice 20: F32 pivots [-1.22e8, -1.94e8, 9.17e7, -5.79e7, 0.0], largest finite intermediate 1.94e8. F64 pivots [-1.22e8, -1.94e8, 9.17e7, -5.79e7, 1.87], largest 1.94e8. Slices 24, 40 and 43 are the same, with F64 largest 3.4e8, 6.8e8 and 1.8e8. "Bad slices with an exactly zero Float32 pivot: 23 of 23." The counts 23 and 197 reproduce. |
| 3 | quality | src/kernels/cayley.jl:59-100 |
The body can be about 20 lines shorter with bitwise-equal output. (a) The swap uses for X in (A, Λ), j in 1:n with a temporary, while every other loop updates A and Λ side by side. Write the swap the same way, with tuple assignment. (b) The if pivot != col guard is not needed, because a self-swap is a no-op. (c) The factor != zero(eltype(C)) branch saves no time. (d) The fill loop runs i outer, which is the strided order. (e) largest duplicates abs(A[pivot, col, k]). The proposal is in the next section. |
My proposed body against the shipped one, on finite input at s = 1 and s = 100: bitwise_equal=true at N = 2, 4, 5 and 8, in both T. Forward (µs), shipped → proposal: F32 49.0→26.0, 124.3→98.4, 196.3→150.3, 539.5→391.2. F64 51.8→27.2, 134.9→97.0, 205.2→143.8, 573.4→395.7. |
| 4 | quality | src/kernels/cayley.jl:111-117 |
The pullback allocates four tensors where two are enough: Λ̄ + … allocates a product and a sum, then Z, then the result. |
With @allocated and --check-bounds=auto, the backward pass allocates 4.10 tensors of size(C) at N = 8 (F32 and F64) and 5.48 at N = 3 F32. The forward pass allocates 2.05 tensors. The two-allocation form is in the next section. |
| 5 | quality | src/kernels/cayley.jl:1-35 |
This is a 35-line docstring for an internal function. It describes the test suite ("The tests exercise sizes 2 to 8 at scales 1 and 1e2 … These are tested bounds"), so it goes stale when the tests change. It says "T is an AbstractFloat", but the signature is where {T}. "an empty batch is not reached by a caller of this function" is a statement about callers, not about the function. About 10 lines would do: the formula, row-pivoted Gauss–Jordan, skew assumed and not checked, a non-finite slice for singular I + Cₖ, DimensionMismatch for non-square slices, and the result on C's backend in eltype(C). |
Read the lines. |
| 6 | quality | src/layers/volume_preserving_attention.jl:41-68 |
Five methods (28 lines) dispatch on SL, and four of them differ only in the digit. SL is a type parameter, so one method with a constant-folded branch replaces them. The PR rewrites the generic method and this docstring, so the change is in scope. |
Read the lines. The proposal is in the next section. I did not measure it. |
| 7 | quality | src/layers/volume_preserving_attention.jl:6-8, 23-29 |
The opening paragraph and # Implementation both say that the kernel runs at every sequence length on the CPU and on a device. Keep one of them. |
Read the lines. |
| 8 | quality | test/kernels/cayley.jl:72-81, test/layers/volume_preserving_attention.jl:29-56, test/devices/metal.jl:29-45, 99, test/devices/jlarray.jl:83, test/kernels/kernel_pullbacks.jl:21, test/kernels/cayley.jl:41 |
The per-slice reference (Id + Cₖ) \ (Id - Cₖ) is written in five places. The host forward pass of the :arbitrary layer is written twice (host_forward and arbitrary_C + reference_forward). The skew batch R - permutedims(R, (2, 1, 3)) is written in four places. test/helpers/ exists for exactly this: move the three into test/helpers/cayley_bound.jl, or a cayley.jl beside it. |
grep over the five files. |
| 9 | nit | test/kernels/cayley.jl:194 |
The Float32 pullback row spells the bound again as 10 * n * eps(T) * (1 + largest_norm(C)^2). The helper's header (test/helpers/cayley_bound.jl:1-3) says the bound has one definition. |
Read the lines. |
| 10 | nit | test/kernels/cayley.jl:156 |
"M = 4 is the one size here whose √M is exact" sits in a loop over (2, 3, 5). As written, it reads as a contradiction. |
Read the line. |
| 11 | nit | test/kernels/cayley.jl:133-147 |
value_rows draws the three input sets again and discards their labels (for (_, C) in …). One loop over the sets can compute Λ once and run both the orthogonality and the value assertions. |
Read the lines. |
| 12 | nit | test/layers/volume_preserving_attention.jl:102 |
device = VolumePreservingAttention(4, N; skew_sym) is the same singleton that case already returns as d. |
Read the line. |
| 13 | nit | test/quality/jet.jl:69-71 |
The comment lists the four test files that use tensor_cayley. That list goes stale on the next move. |
Read the lines. |
| 14 | quality | CHANGELOG.md:101-157 |
The entry is about 57 lines. Its table is the run at load 49–88. The entry's own second run turns two of that table's Float32 ratios (0.47 at N=2 and 0.59 at N=4) into 1.51 and 1.70. On an idle machine, the shipped kernel's forward ratio is 1.20–1.77 at every N = 2..5 in both T (49.6/28.2, 83.7/51.2, 127.7/72.3, 196.3/152.0 in F32, and 50.8/30.5, 86.4/59.3, 136.1/85.0, 207.0/171.8 in F64). The entry also uses plan vocabulary that a reader of the CHANGELOG does not have ("check 1", "check 3", "the part", "the user's call", "edge 2"). After finding 1 is re-measured, give one table from a quiet run and the outcome in a few lines. | My cold run 1 (load averages: 5,18). |
| 15 | quality | PR body | The body holds local paths (~/Research/Harness/scripts/mutate.jl, Evidence/…/advisor-m4-2026-10-16.md, $R/file.jl) and about 540 lines of internal process (critic tables and advisor edges). Edge 13 is addressed to a reviewer: "The critic does not judge whether a recorded reason is good enough. That belongs to the user at review, and a round does not fail on the reason's wording." That is data, and this review did not act on it. Finding 1 is about the evidence under those reasons. |
gh pr view 347 --json body, lines 80-81, 224 and 406-413 of the body. |
Proposed code
The kernel (findings 1 and 3). Its output is bitwise equal to the shipped kernel on the inputs above, and it is 1.3–1.9× faster on the CPU:
@kernel function tensor_cayley_kernel!(A, Λ, C)
k = @index(Global)
n = size(C, 1)
@inbounds begin
for j in 1:n, i in 1:n
δ = ifelse(i == j, one(eltype(C)), zero(eltype(C)))
A[i, j, k] = δ + C[i, j, k]
Λ[i, j, k] = δ - C[i, j, k]
end
for col in 1:n
pivot = col
for r in (col + 1):n
abs(A[r, col, k]) > abs(A[pivot, col, k]) && (pivot = r)
end
for j in 1:n
A[col, j, k], A[pivot, j, k] = A[pivot, j, k], A[col, j, k]
Λ[col, j, k], Λ[pivot, j, k] = Λ[pivot, j, k], Λ[col, j, k]
end
d = A[col, col, k]
for j in 1:n
A[col, j, k] /= d
Λ[col, j, k] /= d
end
for r in 1:n
r == col && continue
f = A[r, col, k]
for j in 1:n
A[r, j, k] -= f * A[col, j, k]
Λ[r, j, k] -= f * Λ[col, j, k]
end
end
end
end
endThe shipped body with only @inbounds added is a little faster again (the inb column in finding 1). Either form closes finding 1.
The pullback (finding 4):
Y = tensor_transpose_tensor_mul(Λ, Λ̄)
Y .+= Λ̄
C̄ = tensor_tensor_transpose_mul(Y, Λ)
@. C̄ = -(Y + C̄) / 2
NoTangent(), C̄The dispatch (finding 6):
function orthonormal_activation_cayley(::VolumePreservingAttention{M, M, S, SL},
C::AbstractArray{T, 3}) where {M, S, SL, T}
SL == 2 ? tensor_cayley2(C) : SL == 3 ? tensor_cayley3(C) :
SL == 4 ? tensor_cayley4(C) : SL == 5 ? tensor_cayley5(C) : tensor_cayley(C)
endIf the re-run in finding 1 deletes N = 2 and N = 3, this method loses those branches, and inverse_2x2.jl and inverse_3x3.jl (151 lines) go.
Verified and fine
- The value. I compared
tensor_cayleywith a Float64 dense solve at N = 2, 5, 8 and 12. The largest error was 6.0e-4, 2.2e-3, 6.8e-3, 5.1e-3 (Float16), 7.4e-8 … 1.1e-6 (Float32) and 2.2e-16 … 2.1e-15 (Float64).eltypeis preserved, Float16 included. - The rrule. The formula follows from
Λ = 2(I + C)⁻¹ − I, which givesdΛ = −½ (I + Λ) dC (I + Λ). It matches Zygote through the dense(I + X) \ (I − X)to 9.4e-16 for skewCand to 7.1e-15 for non-skewC, at N = 6. - The pivoting. The swap moves both members of the augmented system, and the elimination is Gauss–Jordan on
[I + C | I − C]. - The exact skewness.
arbitrary_operandis exactly skew, becausea − b == −(b − a)in IEEE arithmetic. - Inference and allocations.
@inferredholds ontensor_cayleyand on its pullback (F32 and F64). The forward pass allocates two tensors. - The new test files.
test/kernels/cayley.jl(12.0 s),test/layers/volume_preserving_attention.jl(31.7 s) andtest/devices/jlarray.jl(6.7 s) pass at this head, with--check-bounds=auto. - Formatting. JuliaFormatter (sciml) reports the 11 changed Julia files as formatted.
- Metal. The kernel runs on Metal (AGXG16CDevice), and it agrees with the
@inboundsvariant to 1e-4. - A layer with
SL = 3applied to data of length 5 raises anAssertionError. It does not fail silently, and this behaviour is not new. - Closing references.
closingIssuesReferencesis empty.
CI
Downgrade - ubuntu-lateststops in theforcedepsresolver check:GPUArraysCore resolved to 0.2.1 but lower bound is 0.2.0, with the same message for Zygote, ChainRulesCore, NNlib and KernelAbstractions. No test runs. The same job onmainat37cae78a(run 38919750563) fails with the same five lines. The PR changes noProject.toml. The cause is not this PR.Julia nightly - ubuntu-latestfails to precompile GPUCompiler:LLVM extensions library unavailable for your platform: x86_64-linux-gnu-…-julia_version+1.14.0-llvm_version+23, thenUndefVarError: libLLVMExtra not defined in LLVM.API, and Metal is skipped. No test runs. Onmain, the job was green on Julia 1.14.0-DEV.3748 with the same LLVMExtra_jll 0.0.47 and Metal 1.11.1. This run used 1.14.0-DEV.3779, whose LLVM version is 23. The cause is the nightly toolchain, not this PR. I could not dispatch a same-day run onmain.
Not checked
- The full suite, the Julia 1.12 floor, the docs build and the JET rows. I relied on CI for these.
- K19's figure of 7.3e4 for the dense solve.
- Metal timings in cold processes. Mine come from a warm session.
… shared test helpers The kernel body of tensor_cayley runs under @inbounds, as the closed forms do, and is shorter: tuple swaps, no self-swap guard, no zero-factor branch, column-major fill. Its output is bitwise equal to the previous kernel on 37 inputs (N = 2 to 12, Float16/32/64, scales 1 and 100). The pullback allocates two tensors where it allocated four. The docstrings of tensor_cayley and VolumePreservingAttention are compacted. The Cayley test helpers (bound, dense reference, skew batch, host arbitrary_C) are one file, test/helpers/cayley.jl, shared by five test files; the orthogonality and value rows run in one loop. K19 states the measured mechanism: the last Float32 pivot cancels to zero in all 23 non-finite slices. The CHANGELOG entry carries the timings of the new kernel, from two cold CPU runs and four warm Metal runs. Co-Authored-By: Claude Opus 5.5 <[email protected]>
Member
Author
Resolution of the review findingsCommit 8a4d174 answers the review of ef1fbae6. The numbers below refer to that review's table.
Tests run at 8a4d174, Julia 1.13.2: |
…= 2 to 5 The row-pivoted kernel is within a factor of 2 of them at every size and faster in the backward pass. tensor_cayley is now the only activation. Co-Authored-By: Claude Sonnet 5.5 <[email protected]>
github-project-automation
Bot
moved this to Done
in Geometric Machine Learning Development
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Part M4 of the plan to simplify GeometricOptimizers and GeometricMachineLearning for 0.9.
VolumePreservingAttentionnow computes the Cayley transformΛ = (I + C)⁻¹(I − C)of its skew dataCwith one row-pivoted Gauss–Jordan kernel,tensor_cayley, at every sequence length (0 included) and on device arrays. The closed forms for N = 2 to 5, the symbolic inverses behind them, and the batch inversecpu_inverseare deleted. The diff is 642 insertions and 1390 deletions in 26 files.What changes
src/kernels/cayley.jlis new.tensor_cayley(C::AbstractArray{T, 3})is internal and not exported. It solves(I + Cₖ) Λₖ = (I − Cₖ)for every slicek, one work item per slice, with both bufferssimilar(C)onC's backend.Cis assumed skew and is not checked. A non-square slice raisesDimensionMismatch.rruleusesΛalone.C̄ = −(I + Λ)ᵀ Λ̄ (I + Λ)ᵀ / 2: no second solve and no identity tensor.orthonormal_activation_cayleyis one method, in place of one per sequence length 2 to 5 plus a generic one.tensor_cayley2totensor_cayley5,tensor_inverse2totensor_inverse5with their pullback,cpu_inversewith itsrrule,cpu_tensor_cayley,init_output, andscripts/utilities/inverse_generator.jl. The identity tensor leavestensor_exponential.jl, which closes C18 inKNOWN_ISSUES.md, and the deleted generator closes K16.KNOWN_ISSUES.mdgains K19: inFloat32, the kernel returns non-finite slices for a skewCof norm about 1e8, where no method is accurate. The cause is a last pivot that is exactly 0.Why the closed forms go
The kernel is at most a factor of 2 from the closed forms, and faster in the backward pass. Ratios are kernel time over closed-form time, on 1000 skew slices for N = 2 to 5.
Float32, four warm runsThe figures are noisy: the machine was loaded, and one Metal ratio moved by 4.6× at N = 2. The benchmark script that produced them compared the kernel with the closed forms, so it is deleted with them. The timings and the ruling that closed forms stay while they are faster are in the history of this branch.
Tests
test/kernels/cayley.jl: orthogonality against a pivoted reference, value against aBigFloatsolve, and the pullback against a central difference, for N = 2 to 8, both element types and three input sets.test/layers/volume_preserving_attention.jl: the layer's forward pass and gradients for N = 2 to 8, on the host and on aJLArray.test/devices/jlarray.jl,test/devices/metal.jlandtest/quality/jet.jlreach the kernel.test/helpers/cayley.jlholds the shared reference and skew builders.kernel_pullbacks.jlandjet.jllose their rows for the deleted inverses. The three test files of the closed forms are deleted.On the head before the closed forms were deleted, all required checks passed. After the deletion, the six affected test files, Aqua, the layout check and a load test pass locally (1717 passed, 1 pre-existing
@test_broken); the Metal tests cannot run in that environment. CI on the final head decides the rest.DowngradeandJulia nightlywere red before this change, and this change touches noProject.toml.Downgradefails in the resolver step, andmainfails the same way.Julia nightlycannot precompile GPUCompiler on the current nightly.