From 76cee6214abd726dca495f8ea45016fd74a2c4ab Mon Sep 17 00:00:00 2001 From: Michael Kraus Date: Fri, 16 Oct 2026 22:57:10 +0200 Subject: [PATCH 01/10] GML part M4: one row-pivoted Cayley kernel, before the closed forms are decided MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `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`. --- docs/src/layers/attention_layer.md | 1 + scripts/utilities/cayley_timings.jl | 97 +++++++++ src/GeometricMachineLearning.jl | 1 + src/kernels/cayley.jl | 117 +++++++++++ src/layers/volume_preserving_attention.jl | 2 +- test/devices/jlarray.jl | 29 ++- test/devices/metal.jl | 53 ++++- test/kernels/cayley.jl | 227 +++++++++++++++++++++ test/kernels/kernel_pullbacks.jl | 12 +- test/layers/volume_preserving_attention.jl | 156 ++++++++++++++ test/quality/jet.jl | 13 ++ test/runtests.jl | 2 + 12 files changed, 703 insertions(+), 7 deletions(-) create mode 100644 scripts/utilities/cayley_timings.jl create mode 100644 src/kernels/cayley.jl create mode 100644 test/kernels/cayley.jl create mode 100644 test/layers/volume_preserving_attention.jl diff --git a/docs/src/layers/attention_layer.md b/docs/src/layers/attention_layer.md index 96a2ea508..abef60d6f 100644 --- a/docs/src/layers/attention_layer.md +++ b/docs/src/layers/attention_layer.md @@ -258,6 +258,7 @@ Attention was used before [the transformer](@ref "Standard Transformer") was int ```@docs GeometricMachineLearning.tensor_mat_skew_sym_assign +GeometricMachineLearning.tensor_cayley VolumePreservingAttention ``` diff --git a/scripts/utilities/cayley_timings.jl b/scripts/utilities/cayley_timings.jl new file mode 100644 index 000000000..86b236d7a --- /dev/null +++ b/scripts/utilities/cayley_timings.jl @@ -0,0 +1,97 @@ +# The forward and the backward time of the Cayley transform of a batch of 1000 skew slices at scale +# 1, `tensor_cayley`, against the closed forms it would replace: `tensor_cayleyN` for N = 2 to 5 and +# `cpu_tensor_cayley` for N = 6 to 8, each where the tree still defines it. BLAS is pinned to one +# thread, and every variant, every N and every element type runs in a fresh process; a process makes +# one warm-up call it does not time, then takes the median of five timed calls. The parent prints the +# ratio of the kernel to the closed form, forward and backward. +# +# The closed forms that the kernel replaces are deleted by the outcome this table decides, so the +# figures come from a branch commit that holds both the kernel and every closed form, and the +# `CHANGELOG.md` entry of that outcome cites the hash of that commit. After the deletions the script +# times only the variants still defined (`isdefined`). +# +# Run from the repository root, with the scripts environment; one process of one variant is the +# second form, which the first starts: +# julia --startup-file=no --project=scripts scripts/utilities/cayley_timings.jl +# julia --startup-file=no --project=scripts scripts/utilities/cayley_timings.jl tensor_cayley 4 Float32 +# +# Metal has no reference for N = 6 to 8, because `cpu_tensor_cayley` calls `inv` inside a kernel, and +# Metal is not a dependency of `scripts/Project.toml`. The Metal column runs `measure` through +# Kaimon, in a scratch environment that develops the worktree and adds Metal, as +# +# julia> measure("tensor_cayley", 6, Float32, Metal.@sync) +# +# because a Metal launch returns before the kernel has run. That column prints the kernel's own time +# and no ratio. + +using GeometricMachineLearning +using LinearAlgebra: BLAS +using Printf +import Random +import Zygote + +const KERNEL = "tensor_cayley" +const SLICES = 1000 +const RUNS = 5 + +median(v) = sort(v)[(length(v) + 1) ÷ 2] + +# One untimed warm-up call, then `RUNS` timed calls, in µs. `sync` wraps a timed call and is the +# identity except on Metal. +function times(f, sync = identity) + f() + median([(@elapsed sync(f())) * 1e6 for _ in 1:RUNS]) +end + +function skew_batch(::Type{T}, n) where {T} + R = randn(T, n, n, SLICES) + R - permutedims(R, (2, 1, 3)) +end + +# The forward and the backward time of one variant in this process, in µs. +function measure(variant, n, ::Type{T}, sync = identity) where {T} + BLAS.set_num_threads(1) + Random.seed!(0) + + C = skew_batch(T, n) + f = getfield(GeometricMachineLearning, Symbol(variant)) + + forward = times(() -> f(C), sync) + Λ, pullback = Zygote.pullback(f, C) + Λ̄ = rand(T, size(Λ)) + backward = times(() -> pullback(Λ̄), sync) + + forward, backward +end + +# The closed form of the transform of an `n × n` slice, where the tree still defines it. +function reference(n) + n <= 5 ? "tensor_cayley$n" : "cpu_tensor_cayley" +end + +function one_process(variant, n, T) + cmd = `$(Base.julia_cmd()) --startup-file=no --project=$(Base.active_project()) $(@__FILE__) $variant $n $T` + parse.(Float64, split(readchomp(cmd))) +end + +if length(ARGS) == 3 + println(join(measure(ARGS[1], parse(Int, ARGS[2]), getfield(Base, Symbol(ARGS[3]))), " ")) +else + @printf("%-17s %-2s %-8s %11s %11s %s\n", + "variant", "N", "T", "forward µs", "backward µs", + "ratio kernel/closed form") + for T in (Float32, Float64), n in 2:8 + + kernel = one_process(KERNEL, n, T) + name = reference(n) + if isdefined(GeometricMachineLearning, Symbol(name)) + closed = one_process(name, n, T) + @printf("%-17s %-2d %-8s %11.1f %11.1f forward %.2f, backward %.2f\n", + KERNEL, n, T, kernel[1], kernel[2], kernel[1] / closed[1], + kernel[2] / closed[2]) + else + @printf("%-17s %-2d %-8s %11.1f %11.1f no closed form\n", + KERNEL, n, T, kernel[1], kernel[2]) + end + end +end diff --git a/src/GeometricMachineLearning.jl b/src/GeometricMachineLearning.jl index 3c1c850f7..d6fb9af43 100644 --- a/src/GeometricMachineLearning.jl +++ b/src/GeometricMachineLearning.jl @@ -127,6 +127,7 @@ include("kernels/tensor_tensor_mul.jl") include("kernels/tensor_transpose_tensor_mul.jl") include("kernels/tensor_tensor_transpose_mul.jl") include("kernels/mul.jl") +include("kernels/cayley.jl") include("kernels/exponentials/tensor_exponential.jl") include("kernels/inverses/cpu_inverse.jl") include("kernels/inverses/tensor_inverse_pullback.jl") diff --git a/src/kernels/cayley.jl b/src/kernels/cayley.jl new file mode 100644 index 000000000..6d4f79fa3 --- /dev/null +++ b/src/kernels/cayley.jl @@ -0,0 +1,117 @@ +@doc raw""" + tensor_cayley(C::AbstractArray{T, 3}) + +Compute the Cayley transform of every slice of the three-tensor `C`: the slice ``k`` of the result +is + +```math + \Lambda_k = (\mathbb{I} + C_k)^{-1}(\mathbb{I} - C_k), +``` + +computed by row-pivoted Gauss–Jordan elimination on the pair ``(\mathbb{I} + C_k, +\mathbb{I} - C_k)``, which leaves ``\Lambda_k`` in the second member. This is the transform of +[`VolumePreservingAttention`](@ref), which forms ``C`` skew-symmetric. + +`C` is assumed skew-symmetric, and this is not checked: a check would cost a reduction per slice on +a device. The result is the Cayley transform of any `C` whose ``\mathbb{I} + C_k`` is nonsingular; +for a skew `C` its orthogonality error ``\|\Lambda_k^T\Lambda_k - \mathbb{I}\|_\infty`` is the error +of a row-pivoted elimination — at most four times that of a dense `(I + C_k) \ (I - C_k)`, and at +most ``10n\varepsilon(1 + \|C_k\|_2^2)`` for the element type's machine epsilon ``\varepsilon``, +where `n = size(C, 1)`. A skew-symmetric ``C_k`` has the positive definite symmetric part +``\mathbb{I}``, which is what bounds the pivots. + +The element type `T` is an `AbstractFloat`. Every buffer comes from `similar(C)`, so the result has +the element type `T` and lives on the backend of `C`. The loop bound is read from `size(C, 1)` at +run time, so a slice of any size is solved by the same kernel. + +A slice whose ``\mathbb{I} + C_k`` is singular gives non-finite entries in the slice of that `k` and +no error, as the closed forms do, which divide by the determinant of the slice; a `NaN` or an `Inf` +in `C` spreads within its own slice. A slice that is not square raises a `DimensionMismatch`, and +an empty batch is not reached by a caller of this function. +""" +function tensor_cayley(C::AbstractArray{T, 3}) where {T} + n = size(C, 1) + size(C, 2) == n || + throw(DimensionMismatch("a slice of C is $(n)×$(size(C, 2)); the slices must be square")) + + backend = networkbackend(C) + Λ = similar(C) + A = similar(C) + + cayley! = tensor_cayley_kernel!(backend) + cayley!(A, Λ, C, ndrange = size(C, 3)) + + Λ +end + +# `(I + Cₖ) Λₖ = (I - Cₖ)` for every slice at once, one work item per slice: `A` is filled with +# `I + C`, `Λ` with `I - C`, and the row operations that turn `A` into `I` are also applied to `Λ`. +# The pivot of a column is the row of the largest magnitude in it, and the swap moves both members: +# the two are one augmented system, and a swap on one of them alone solves a permuted system. +@kernel function tensor_cayley_kernel!(A, Λ, C) + k = @index(Global) + n = size(C, 1) + + for i in 1:n, j in 1:n + + δ = 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 + largest = abs(A[col, col, k]) + for r in (col + 1):n + candidate = abs(A[r, col, k]) + if candidate > largest + largest = candidate + pivot = r + end + end + + if pivot != col + for j in 1:n + a = A[col, j, k] + A[col, j, k] = A[pivot, j, k] + A[pivot, j, k] = a + b = Λ[col, j, k] + Λ[col, j, k] = Λ[pivot, j, k] + Λ[pivot, j, k] = b + end + end + + diagonal = A[col, col, k] + for j in 1:n + A[col, j, k] = A[col, j, k] / diagonal + Λ[col, j, k] = Λ[col, j, k] / diagonal + end + + for r in 1:n + factor = A[r, col, k] + if r != col && factor != zero(eltype(C)) + for j in 1:n + A[r, j, k] = A[r, j, k] - factor * A[col, j, k] + Λ[r, j, k] = Λ[r, j, k] - factor * Λ[col, j, k] + end + end + end + end + + nothing +end + +# `Λ = (I + C)⁻¹(I - C)` gives `(I + Λ) = 2(I + C)⁻¹`, so the cotangent of `C` under the Frobenius +# product is `C̄ = -(I + Λ)ᵀ Λ̄ (I + Λ)ᵀ / 2`, from `Λ` alone: no solve, and no identity tensor. +function ChainRulesCore.rrule(::typeof(tensor_cayley), C::AbstractArray{T, 3}) where {T} + Λ = tensor_cayley(C) + + function tensor_cayley_pullback(Λ̄) + Λ̄ = _seed(T, Λ̄) + Λ̄ isa AbstractZero && return NoTangent(), ZeroTangent() + Y = Λ̄ + tensor_transpose_tensor_mul(Λ, Λ̄) + NoTangent(), -(Y + tensor_tensor_transpose_mul(Y, Λ)) / 2 + end + + Λ, tensor_cayley_pullback +end diff --git a/src/layers/volume_preserving_attention.jl b/src/layers/volume_preserving_attention.jl index 2b7c4901f..e862928f4 100644 --- a/src/layers/volume_preserving_attention.jl +++ b/src/layers/volume_preserving_attention.jl @@ -37,7 +37,7 @@ end function orthonormal_activation_cayley(::VolumePreservingAttention, A::AbstractArray{ T, 3}) where {T} - cpu_tensor_cayley(A) + tensor_cayley(A) end function orthonormal_activation_cayley( diff --git a/test/devices/jlarray.jl b/test/devices/jlarray.jl index ac08faf71..eeb9cf047 100644 --- a/test/devices/jlarray.jl +++ b/test/devices/jlarray.jl @@ -1,12 +1,13 @@ -# The products of `src/kernels/` on a device array that has no BLAS and no NNlib extension: a -# `JLArray` with scalar indexing off, so that a product that reads one element at a time raises. +# The kernels of `src/kernels/` on a device array that has no BLAS and no NNlib extension: a +# `JLArray` with scalar indexing off, so that a kernel that reads one element at a time raises. # Each result must stay on `JLBackend()` and match the host. using GeometricMachineLearning -using GeometricMachineLearning: _mul, tensor_tensor_mul, assign_q_and_p +using GeometricMachineLearning: _mul, tensor_tensor_mul, tensor_cayley, assign_q_and_p using ChainRulesCore: rrule using GPUArraysCore: allowscalar using JLArrays: JLArray, JLBackend +using LinearAlgebra: I, opnorm import KernelAbstractions using Zygote: gradient using Test @@ -59,3 +60,25 @@ end end end end + +# The Cayley kernel's own device rows: it launches where scalar indexing is off, its result stays on +# the device, and it matches the per-slice solve on the host at the bound of its orthogonality rows. +@testset "tensor_cayley and its pullback on a JLArray" begin + for T in (Float32, Float64), n in (2, 3, 8) + + R = rand(T, n, n, 4) + C = R - permutedims(R, (2, 1, 3)) + Λ, pullback = rrule(tensor_cayley, JLArray(C)) + Λ̄ = rand(T, n, n, 4) + _, C̄ = pullback(JLArray(Λ̄)) + @test ondevice(Λ) && ondevice(C̄) + + Id = Matrix{T}(I, n, n) + reference = cat(((Id + C[:, :, k]) \ (Id - C[:, :, k]) for k in axes(C, 3))...; dims = 3) + bound = 10 * n * eps(T) * (1 + maximum(opnorm(C[:, :, k], 2)^2 for k in axes(C, 3))) + @test maximum(abs, Array(Λ) - reference) <= bound + + _, C̄_host = rrule(tensor_cayley, C)[2](Λ̄) + @test maximum(abs, Array(C̄) - C̄_host) <= bound * maximum(abs, C̄_host) + end +end diff --git a/test/devices/metal.jl b/test/devices/metal.jl index a12234279..127d0f215 100644 --- a/test/devices/metal.jl +++ b/test/devices/metal.jl @@ -13,12 +13,34 @@ # that reads an array one element at a time raises instead of answering slowly. using GeometricMachineLearning -using GeometricMachineLearning: _mul, tensor_tensor_mul +using GeometricMachineLearning: _mul, tensor_tensor_mul, tensor_cayley import KernelAbstractions +using LinearAlgebra: I, dot, opnorm using Metal using Test import Random +# The layer's forward pass on the host, slice by slice: the same scalar products the kernel computes +# for `C`, a dense solve of the Cayley transform, and the product with the data. It returns the +# slices of `C` as well, because the bound of the kernel's orthogonality rows is written in them. +function host_forward(x::AbstractArray{T, 3}, A::AbstractMatrix) where {T} + M = size(x, 1) + out = similar(x) + C = zeros(T, M, M, size(x, 3)) + for k in axes(x, 3) + X = Matrix(view(x, :, :, k)) + L = zeros(T, M, M) + for i in 1:M, j in 1:(i - 1) + + L[i, j] = dot(view(X, :, i), A * view(X, :, j)) + end + C[:, :, k] = (L - L') / T(√M) + Id = Matrix{T}(I, M, M) + out[:, :, k] = X * ((Id + C[:, :, k]) \ (Id - C[:, :, k])) + end + out, C +end + if Metal.functional() Metal.allowscalar(false) Random.seed!(123) @@ -60,6 +82,35 @@ if Metal.functional() @test Array(device) ≈ host end end + + # The Cayley kernel and the layer that uses it, at every sequence length: each result stays on + # the device and matches a host computation of the same transform, at the bound of the kernel's + # orthogonality rows. + @testset "tensor_cayley and VolumePreservingAttention on Metal" begin + for n in 2:8 + R = rand(T, n, n, 8) + C = R - permutedims(R, (2, 1, 3)) + Λ = tensor_cayley(MtlArray(C)) + @test KernelAbstractions.get_backend(Λ) == backend + Id = Matrix{T}(I, n, n) + reference = cat(((Id + C[:, :, k]) \ (Id - C[:, :, k]) for k in axes(C, 3))...; + dims = 3) + bound = 10 * n * eps(T) * + (1 + maximum(opnorm(C[:, :, k], 2)^2 for k in axes(C, 3))) + @test maximum(abs, Array(Λ) - reference) <= bound + + d = VolumePreservingAttention(n, n) + ps = NeuralNetwork(Chain(d), CPU(), T).params.L1 + x = rand(T, n, n, 8) + y = d(MtlArray(x), (A = MtlArray(ps.A),)) + expected, C_layer = host_forward(x, ps.A) + layer_bound = 10 * n * eps(T) * + (1 + maximum(opnorm(C_layer[:, :, k], 2)^2 for k in axes(x, 3))) * + maximum(abs, x) + @test KernelAbstractions.get_backend(y) == backend + @test maximum(abs, Array(y) - expected) <= layer_bound + end + end else @test_skip Metal.functional() end diff --git a/test/kernels/cayley.jl b/test/kernels/cayley.jl new file mode 100644 index 000000000..10b374882 --- /dev/null +++ b/test/kernels/cayley.jl @@ -0,0 +1,227 @@ +# `tensor_cayley`: that its rows are skew on every input the layer sends it, that its orthogonality +# error is the error of a row-pivoted elimination, that its value is the value of the transform, and +# that its pullback is the derivative of that value. +# +# The three input sets are the ones the layer produces: a hand-built skew tensor, the `:arbitrary` +# path and the `:skew_sym` path, both of them through the layer's own functions. The reference is the +# per-slice solve `(I + Cₖ) \ (I - Cₖ)` in the working precision, written here, because the closed +# forms are what `tensor_cayley` replaces. + +using GeometricMachineLearning +using GeometricMachineLearning: tensor_cayley, tensor_mat_skew_sym_assign, + tensor_transpose_tensor_mul, _mul +using ChainRulesCore: AbstractZero, ZeroTangent, rrule +using LinearAlgebra: BLAS, I, opnorm +using Test +import Random + +Random.seed!(1234) + +# The references below are solves of 8×8 systems, thousands of them, and threaded LAPACK spends more +# time waking its threads than solving: with BLAS on this machine's twelve threads the orthogonality +# rows over their 84 batches take about two minutes, and 2.6 s on one thread, where the value rows +# take 10.9 s. The file pins one thread for its own work and restores the count it found at the end, +# so no later testset inherits it. +const BLAS_THREADS = BLAS.get_num_threads() +BLAS.set_num_threads(1) + +const SLICES = 1000 +const SCALES = (1, 1e2) + +function identity_tensor(::Type{T}, n::Int, k::Int) where {T} + E = zeros(T, n, n, k) + for i in 1:n + E[i, i, :] .= one(T) + end + E +end + +# The largest `‖Cₖ‖₂` over the slices of `C`, as a `Float64`. +function largest_norm(C::AbstractArray) + maximum(Float64(opnorm(@view(C[:, :, k]), 2)) for k in axes(C, 3)) +end + +# The tensor scaled so that its largest `‖Cₖ‖₂` is `s`. +scale_to(C::AbstractArray, s) = C * (s / largest_norm(C)) + +# (i) A hand-built skew tensor, `R - Rᵀ` for `R = s * randn`, which is skew to rounding. +function skew_tensor(::Type{T}, n::Int, s, k::Int = SLICES) where {T} + R = T(s) * randn(T, n, n, k) + R - permutedims(R, (2, 1, 3)) +end + +# (ii) The `:arbitrary` path of the layer: `tensor_mat_skew_sym_assign(x, A) / √M`, then +# `x_interim - permutedims(x_interim, (2, 1, 3))`. That path is linear in `A`, so the whole tensor is +# scaled by one factor on `A` to put its largest `‖Cₖ‖₂` at `s`. +function arbitrary_tensor(::Type{T}, n::Int, s) where {T} + x = randn(T, n, n, SLICES) + A = randn(T, n, n) + scale_to(arbitrary_path(x, A), s) +end + +function arbitrary_path(x::AbstractArray{T, 3}, A::AbstractMatrix) where {T} + x_interim = tensor_mat_skew_sym_assign(x, A) / T(√size(x, 1)) + x_interim - permutedims(x_interim, (2, 1, 3)) +end + +# (iii) The `:skew_sym` path of the layer: `Zᵀ A Z` for a `SkewSymMatrix` `A`. That is quadratic in +# `Z`, so `Z` is scaled by the square root of the factor that puts the largest `‖Cₖ‖₂` at `s`. +function skew_sym_tensor(::Type{T}, n::Int, s) where {T} + x = randn(T, n, n, SLICES) + A = rand(SkewSymMatrix{T}, n) + C = skew_sym_path(x, A) + skew_sym_path(x * T(sqrt(s / largest_norm(C))), A) +end + +function skew_sym_path(x::AbstractArray{T, 3}, A::AbstractMatrix) where {T} + tensor_transpose_tensor_mul(x, _mul(A, x)) +end + +# The per-slice reference of the orthogonality rows and of the value row: the solve +# `(I + Cₖ) \ (I - Cₖ)` in the working precision, and in `BigFloat` for the value row. The two +# agree to the last bits of `T`, so the value row's reference is the transform itself rather than +# another elimination. +function reference_solve(::Type{T}, Ck::AbstractMatrix) where {T} + Id = Matrix{T}(I, size(Ck)) + (Id + Ck) \ (Id - Ck) +end + +function big_reference_solve(::Type{T}, Ck::AbstractMatrix) where {T} + Id = Matrix{BigFloat}(I, size(Ck)) + C = BigFloat.(Ck) + T.((Id + C) \ (Id - C)) +end + +# The same maximum over the pivoted references, which is what the ratio clause compares against. +function reference_orthogonality_error(C::AbstractArray{T, 3}) where {T} + n = size(C, 1) + E = Matrix{T}(I, n, n) + worst = 0.0 + for k in axes(C, 3) + X = reference_solve(T, @view(C[:, :, k])) + worst = max(worst, Float64(maximum(abs, X' * X - E))) + end + worst +end + +# The bounds of the orthogonality rows, per slice: the absolute clause `10 n eps(T) (1 + ‖Cₖ‖₂²)`, +# and the ratio clause's own factor on the reference's error. +function bound_per_slice(::Type{T}, Ck::AbstractMatrix, n::Int) where {T} + 10.0 * n * eps(T) * (1.0 + Float64(opnorm(Ck, 2))^2) +end + +# Every row of the orthogonality check: the ratio against the pivoted reference on all three input +# sets, and the absolute bound on the two exactly skew ones. +function orthogonality_rows(::Type{T}, n::Int, s) where {T} + kernel_error = Ref(0.0) + reference_error = Ref(0.0) + worst_absolute = Ref(0.0) + for (set, C) in (("hand-built", skew_tensor(T, n, s)), + ("arbitrary", arbitrary_tensor(T, n, s)), + ("skew_sym", skew_sym_tensor(T, n, s))) + Λ = tensor_cayley(C) + E = identity_tensor(T, n, size(C, 3)) + ΛᵀΛ = tensor_transpose_tensor_mul(Λ, Λ) + kernel_error[] = Float64(maximum(abs, ΛᵀΛ - E)) + reference_error[] = reference_orthogonality_error(C) + + # The ratio clause holds on every input set; the reference's own error is not zero, so the + # factor is a real bound rather than an equality. + @test kernel_error[] <= 4 * reference_error[] + + if set != "skew_sym" + for k in axes(C, 3) + bound = bound_per_slice(T, @view(C[:, :, k]), n) + error_k = Float64(maximum(abs, (@view ΛᵀΛ[:, :, k]) .- (@view E[:, :, k]))) + worst_absolute[] = max(worst_absolute[], error_k / bound) + end + @test worst_absolute[] <= 1 + end + end +end + +# The value row: per slice, the distance from the `BigFloat` solve, rounded to `T`, under the same +# bound. `Λᵀ` is orthogonal wherever `Λ` is, so the orthogonality rows alone do not pin the value. +function value_rows(::Type{T}, n::Int, s) where {T} + for (set, C) in (("hand-built", skew_tensor(T, n, s)), + ("arbitrary", arbitrary_tensor(T, n, s))) + Λ = tensor_cayley(C) + worst = Ref(0.0) + for k in axes(C, 3) + bound = bound_per_slice(T, @view(C[:, :, k]), n) + error_k = Float64(maximum(abs, (@view Λ[:, :, k]) .- + big_reference_solve(T, @view(C[:, :, k])))) + worst[] = max(worst[], error_k / bound) + end + @test worst[] <= 1 + end +end + +@testset "tensor_cayley: the input of the layer is always skew" begin + for T in (Float32, Float64), n in 2:4 + + x = randn(T, n, n, 7) + A = randn(T, n, n) + # The `:arbitrary` path is exactly skew: one subtraction of a permutation of the same array. + C = arbitrary_path(x, A) + @test iszero(C + permutedims(C, (2, 1, 3))) + # The `:skew_sym` path is skew only to rounding, so nothing is asserted exactly on it. Its + # deviation is far below the orthogonality rows' bound, which is what those rows measure. + C2 = skew_sym_path(x, rand(SkewSymMatrix{T}, n)) + deviation = maximum(abs, C2 + permutedims(C2, (2, 1, 3))) + @test deviation <= 100 * eps(T) * n * maximum(abs, C2) + end +end + +@testset "tensor_cayley: the orthogonality error and the value, N = $n, s = $s" for T in (Float32, Float64), + s in SCALES, n in 2:8 + + orthogonality_rows(T, n, s) + value_rows(T, n, s) +end + +@testset "tensor_cayley: the pullback, N = $n" for T in (Float32, Float64), + n in (2, 3, 5, 8) + + C = skew_tensor(T, n, 1, 7) + Λ̄ = randn(T, n, n, 7) + + if T === Float64 + # The relative error against a central difference of `⟨Λ̄, tensor_cayley(·)⟩` along a random + # direction, with `h = cbrt(eps)`. + C̄ = rrule(tensor_cayley, C)[2](Λ̄)[2] + D = randn(T, size(C)) + h = cbrt(eps(T)) + loss(X) = sum(Λ̄ .* tensor_cayley(X)) + fd = (loss(C + h * D) - loss(C - h * D)) / (2h) + @test abs(sum(C̄ .* D) - fd) / abs(fd) <= 1e-6 + else + # The `Float32` pullback against the `Float64` one, under the bound of the orthogonality + # rows: the cotangent is compared relatively, because its scale is the data's. + C̄32 = rrule(tensor_cayley, C)[2](Λ̄)[2] + C̄64 = rrule(tensor_cayley, Float64.(C))[2](Float64.(Λ̄))[2] + bound = 10 * n * eps(T) * (1 + largest_norm(C)^2) + @test maximum(abs, C̄32 - Float32.(C̄64)) <= bound * maximum(abs, C̄64) + end + + # The seed handling of every kernel rule: a zero seed, and a seed of the other precision. + @test all(c -> c isa AbstractZero, rrule(tensor_cayley, C)[2](ZeroTangent())) + @test eltype(rrule(tensor_cayley, C)[2](Float64.(Λ̄))[2]) === T +end + +@testset "tensor_cayley: the result's element type, backend and shape" begin + for T in (Float32, Float64), n in 2:8 + + C = skew_tensor(T, n, 1) + Λ = @inferred tensor_cayley(C) + @test Λ isa Array{T, 3} + @test size(Λ) == size(C) + @test @inferred(tensor_cayley(C))[1, 1, 1] isa T + end + for T in (Float32, Float64) + @test_throws DimensionMismatch tensor_cayley(rand(T, 3, 4, 2)) + @test_throws DimensionMismatch tensor_cayley(rand(T, 4, 3, 2)) + end +end + +BLAS.set_num_threads(BLAS_THREADS) diff --git a/test/kernels/kernel_pullbacks.jl b/test/kernels/kernel_pullbacks.jl index f4815a587..9351d1b38 100644 --- a/test/kernels/kernel_pullbacks.jl +++ b/test/kernels/kernel_pullbacks.jl @@ -1,7 +1,7 @@ using GeometricMachineLearning: tensor_tensor_mul, tensor_transpose_tensor_mul, tensor_mat_skew_sym_assign, cpu_inverse, tensor_inverse2, tensor_inverse3, tensor_inverse4, tensor_inverse5, - init_output, + init_output, tensor_cayley, _mul using ChainRulesCore: rrule, AbstractZero, NoTangent, ZeroTangent, @thunk using ChainRulesTestUtils @@ -17,6 +17,12 @@ const verbose = false # `n × n` slices that are far from singular, for the inverse rules. invertible(T, n, k) = cat((rand(T, n, n) + n * I for _ in 1:k)...; dims = 3) +# A skew batch for the Cayley rule, whose slices are square. +function skew(T, n, k) + R = rand(T, n, n, k) + R - permutedims(R, (2, 1, 3)) +end + function main(first_dim, second_dim, third_dim, third_tensor_dim) test_rrule(tensor_tensor_mul, rand(first_dim, second_dim, third_tensor_dim), rand(second_dim, third_dim, third_tensor_dim)) @@ -29,6 +35,7 @@ function main(first_dim, second_dim, third_dim, third_tensor_dim) test_rrule(_mul, rand(first_dim, second_dim)', rand(first_dim, third_dim)) test_rrule(cpu_inverse, invertible(Float64, first_dim, third_tensor_dim)) test_rrule(init_output, rand(first_dim, first_dim, third_tensor_dim)) + test_rrule(tensor_cayley, skew(Float64, first_dim, third_tensor_dim)) end const dim_range = 10 @@ -60,7 +67,8 @@ function rules(T) (tensor_inverse3, (invertible(T, 3, 2),)), (tensor_inverse4, (invertible(T, 4, 2),)), (tensor_inverse5, (invertible(T, 5, 2),)), - (init_output, (rand(T, 3, 3, 2),))) + (init_output, (rand(T, 3, 3, 2),)), + (tensor_cayley, (skew(T, 3, 2),))) end @testset "a zero seed gives a zero cotangent for every argument" begin diff --git a/test/layers/volume_preserving_attention.jl b/test/layers/volume_preserving_attention.jl new file mode 100644 index 000000000..d082d7884 --- /dev/null +++ b/test/layers/volume_preserving_attention.jl @@ -0,0 +1,156 @@ +# `VolumePreservingAttention` at every sequence length: the forward pass against an independent host +# computation of the same transform, and the gradient against the gradient on `Array`. The layer's +# own `Array` result is not the reference, because the device and the host run the same kernel. +# +# The host reference builds `C` with plain products and solves `(I + C) \ (I - C)` per slice; the +# tolerance of the forward pass is the bound of the Cayley kernel's value row, `10 n eps(T) (1 + +# ‖Cₖ‖₂²)`, times the size of `Z`, which is what the product that follows the transform scales it by. + +using GeometricMachineLearning +using GeometricMachineLearning: tensor_mat_skew_sym_assign, _mul +using FiniteDifferences: central_fdm, j′vp +using GPUArraysCore: allowscalar +using JLArrays: JLArray, JLBackend +import KernelAbstractions +using LinearAlgebra: I, dot, opnorm +using Zygote: gradient +using Test +import Random + +Random.seed!(1234) +allowscalar(false) + +ondevice(x) = KernelAbstractions.get_backend(x) isa JLBackend + +function cayley_reference(C::AbstractMatrix{T}) where {T} + Id = Matrix{T}(I, size(C)) + (Id + C) \ (Id - C) +end + +# The `:arbitrary` path's `C`, slice by slice: the lower triangle of the scalar products of the +# columns of `Z` weighted by `A`, divided by `√M`, and made skew by subtracting its own transpose — +# the same triangle, the same divisor and the same subtraction the layer makes. +function arbitrary_C(x::AbstractArray{T, 3}, A::AbstractMatrix) where {T} + M, seq = size(x, 1), size(x, 2) + C = zeros(T, seq, seq, size(x, 3)) + for k in axes(x, 3) + X = Matrix(@view x[:, :, k]) + L = zeros(T, seq, seq) + for i in 1:seq, j in 1:(i - 1) + + L[i, j] = dot(@view(X[:, i]), A * @view(X[:, j])) + end + C[:, :, k] = (L - L') / T(√M) + end + C +end + +function skew_sym_C(x::AbstractArray{T, 3}, A) where {T} + C = zeros(T, size(x, 2), size(x, 2), size(x, 3)) + D = Matrix(A) + for k in axes(x, 3) + X = Matrix(@view x[:, :, k]) + C[:, :, k] = X' * (D * X) + end + C +end + +function C_tensor(d, x, ps) + d isa VolumePreservingAttention{<:Any, <:Any, :skew_sym} ? + skew_sym_C(x, ps.A) : arbitrary_C(x, ps.A) +end + +# The layer's forward pass, slice by slice on the host: the same `C` the layer forms, a dense solve +# of the Cayley transform, and the product with the data. +function reference_forward(d::VolumePreservingAttention, x, ps) + C = C_tensor(d, x, ps) + out = similar(x) + for k in axes(x, 3) + out[:, :, k] = Matrix(@view x[:, :, k]) * cayley_reference(@view C[:, :, k]) + end + out +end + +# The bound of the Cayley kernel's value row, scaled by the largest entry of the data: the product +# `Z Λ` multiplies the transform's error by `Z`. +function tolerance(d, x, ps) + T = eltype(x) + C = C_tensor(d, x, ps) + worst = maximum(Float64(opnorm(@view(C[:, :, k]), 2))^2 for k in axes(x, 3)) + 10 * size(x, 2) * eps(T) * (1 + worst) * Float64(maximum(abs, Array(x))) +end + +# The layer, its parameters, and the data of one case, all on the host; the device case converts +# them, as `test/devices/jlarray.jl` does, rather than initialising them on the device. +function case(N, T; skew_sym) + d = VolumePreservingAttention(N, N; skew_sym) + model = Chain(d) + ps = NeuralNetwork(model, CPU(), T).params.L1 + x = randn(T, N, N, 3) + d, ps, x +end + +@testset "VolumePreservingAttention runs at every sequence length, T = $T" for T in (Float32, Float64) + for N in 2:8, skew_sym in (false, true) + + d, ps, x = case(N, T; skew_sym) + device = VolumePreservingAttention(N, N; skew_sym) + ps_device = (A = skew_sym ? SkewSymMatrix(JLArray(ps.A.S), N) : JLArray(ps.A),) + + y_device = device(JLArray(x), ps_device) + @test ondevice(y_device) + + # The forward pass against the independent host computation. + reference = reference_forward(d, x, ps) + @test maximum(abs, Array(y_device) - reference) <= tolerance(d, x, ps) + + # The gradient with respect to the data, on the device against the host. + gx_device = gradient(z -> sum(device(z, ps_device)), JLArray(x))[1] + gx_host = gradient(z -> sum(d(z, ps)), x)[1] + @test ondevice(gx_device) + @test maximum(abs, Array(gx_device) - gx_host) <= + tolerance(d, x, ps) * 10 + + # The gradient with respect to the weight, on the device against the host. + gA_device = gradient(A -> sum(device(JLArray(x), (A = A,))), ps_device.A)[1] + gA_host = gradient(A -> sum(d(x, (A = A,))), ps.A)[1] + if skew_sym + @test ondevice(gA_device.S) + @test maximum(abs, Array(gA_device.S) - gA_host.S) <= tolerance(d, x, ps) * 10 + else + @test ondevice(gA_device) + @test maximum(abs, Array(gA_device) - gA_host) <= tolerance(d, x, ps) * 10 + end + end +end + +@testset "VolumePreservingAttention with SL = 0 runs every sequence length, T = $T" for T in (Float32, + Float64) + for seq_length in (3, 7), skew_sym in (false, true) + + d = VolumePreservingAttention(4; skew_sym) + model = Chain(d) + ps = NeuralNetwork(model, CPU(), T).params.L1 + x = randn(T, 4, seq_length, 3) + + y = d(x, ps) + @test size(y) == size(x) + @test maximum(abs, y - reference_forward(d, x, ps)) <= tolerance(d, x, ps) + + ps_device = (A = skew_sym ? SkewSymMatrix(JLArray(ps.A.S), 4) : JLArray(ps.A),) + y_device = d(JLArray(x), ps_device) + @test ondevice(y_device) + @test maximum(abs, Array(y_device) - reference_forward(d, x, ps)) <= + tolerance(d, x, ps) + end +end + +@testset "the gradient of VolumePreservingAttention at N = 6 against finite differences" begin + T = Float64 + d, ps, x = case(6, T; skew_sym = false) + Λ̄ = randn(T, size(x)) + fd = j′vp(central_fdm(5, 1), z -> d(z, ps), Λ̄, x)[1] + analytic = gradient(z -> sum(Λ̄ .* d(z, ps)), x)[1] + # `1e-8` in `Float64`, as `test/integration/product_gradients.jl` measures the products. + @test isapprox(analytic, fd; rtol = 1e-8) +end diff --git a/test/quality/jet.jl b/test/quality/jet.jl index 4c8edf506..19713b244 100644 --- a/test/quality/jet.jl +++ b/test/quality/jet.jl @@ -66,6 +66,19 @@ const A3{T} = Array{T, 3} zeros(T, 3, 3, 2))) end + # src/kernels/cayley.jl: `tensor_cayley` in test/kernels/cayley.jl, + # test/layers/volume_preserving_attention.jl, test/devices/jlarray.jl and + # test/devices/metal.jl + for T in (Float16, Float32, Float64) + @test isempty(JET.get_reports(JET.report_opt(GML.tensor_cayley, (A3{T},); + target_modules = TARGET))) + @test isempty(kernel_body_reports(GML.tensor_cayley_kernel!, 2, + zeros(T, 3, 3, 2), zeros(T, 3, 3, 2), zeros(T, 3, 3, 2))) + @test isempty(JET.get_reports(JET.report_opt( + rrule(GML.tensor_cayley, zeros(T, 3, 3, 2))[2], (A3{T},); + target_modules = TARGET))) + end + # src/kernels/kernel_ad_routines/: the pullback that launches kernels, at the `Float64` # types of `test_rrule` in test/kernels/kernel_pullbacks.jl B, dC = zeros(3, 3, 2), zeros(3, 3, 2) diff --git a/test/runtests.jl b/test/runtests.jl index ebc25815a..09963a200 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -18,6 +18,7 @@ if "core" in GROUPS @safetestset "A Stiefel network trains on data of the other precision" include("integration/mixed_precision_training.jl") @safetestset "Test parallel inverses" include("kernels/tensor_inverse.jl") @safetestset "Test parallel Cayley" include("kernels/tensor_cayley.jl") + @safetestset "tensor_cayley: skew inputs, the orthogonality bound and the pullback" include("kernels/cayley.jl") @safetestset "Volume-Preserving Transformer (skew-symmetric tests)" include("kernels/test_skew_map.jl") @safetestset "Volume-Preserving Transformer (cayley-transform tests)" include("kernels/inverses/test_cayley_transforms.jl") @safetestset "Gradient Layer" include("layers/gradient_layer_tests.jl") @@ -30,6 +31,7 @@ if "core" in GROUPS @safetestset "parameterlength(::PSDLayer) is exact" include("layers/psd_parameterlength.jl") @safetestset "parameterlength(::GrassmannLayer) is an integer count" include("layers/grassmann_parameterlength.jl") @safetestset "Attention layer #1" include("layers/attention_setup.jl") + @safetestset "VolumePreservingAttention at every sequence length" include("layers/volume_preserving_attention.jl") @safetestset "Test setup of MultiHeadAttention layer Stiefel weights" include("layers/multi_head_attention_stiefel_setup.jl") @safetestset "Test geodesic and Cayley retr for the MultiHeadAttention layer w/ St weights" include("layers/multi_head_attention_stiefel_retraction.jl") @safetestset "Test the correct setup of the various optimizer caches for MultiHeadAttention" include("layers/multi_head_attention_stiefel_optim_cache.jl") From dda31c7ffa6a22ad2620ae154cd2bffe580440fb Mon Sep 17 00:00:00 2001 From: Michael Kraus Date: Fri, 16 Oct 2026 23:22:02 +0200 Subject: [PATCH 02/10] GML part M4: the kernel keeps the closed forms for 2 to 5 and takes over the batch inverse MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The benchmark of check 1 ran at 76cee621, 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. --- CHANGELOG.md | 45 ++++++++++++++- KNOWN_ISSUES.md | 10 ---- scripts/utilities/cayley_timings.jl | 56 ++++++++++--------- scripts/utilities/inverse_generator.jl | 2 +- src/GeometricMachineLearning.jl | 3 +- ...nsor_exponential.jl => identity_tensor.jl} | 4 +- src/kernels/inverses/cpu_inverse.jl | 48 ---------------- src/kernels/inverses/tensor_cayley.jl | 5 -- src/layers/volume_preserving_attention.jl | 14 +++-- .../inverses/test_cayley_transforms.jl | 4 +- test/kernels/kernel_pullbacks.jl | 4 +- test/kernels/tensor_inverse.jl | 20 +------ test/quality/jet.jl | 13 +---- 13 files changed, 90 insertions(+), 138 deletions(-) rename src/kernels/{exponentials/tensor_exponential.jl => identity_tensor.jl} (80%) delete mode 100644 src/kernels/inverses/cpu_inverse.jl diff --git a/CHANGELOG.md b/CHANGELOG.md index 18d295bb5..0d1ecbcce 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -73,7 +73,7 @@ breaking release). - **`assign_q_and_p` splits its input by slicing, and a loss that reads one half no longer raises.** It takes `x[1:N, …]` and `x[N+1:2N, …]` with `N::Integer`, where it took `N::Int`. Its six kernels, its rrule and `_vcat_qp` are gone. The old rrule built a cotangent of `N + 1` rows, `vcat(ZeroTangent(), p̄)`, and raised `DimensionMismatch` when a loss read only one of the two halves. The test is `test/kernels/assign_q_and_p.jl`. -- **The kernel rrules accept a zero seed and a seed of the other float type.** Each remaining kernel rrule (`tensor_tensor_mul`, `tensor_transpose_tensor_mul`, `tensor_mat_skew_sym_assign`, `cpu_inverse`, `tensor_inverse2` to `tensor_inverse5`, and `_mul(::Adjoint, x)`) answers a zero seed (`ZeroTangent`, `NoTangent` or a thunk of one) with `NoTangent()` or `ZeroTangent()`, where it raised a `MethodError`. It converts a seed of the other float type to the element type of the primal: a Float64 seed on a Float32 primal raised a `MethodError`, and `cpu_inverse` returned a Float32 `dA` for a Float64 primal. `init_output` already did this. The style is the same in every rule: eager, with no `@thunk`. The four generated `tensor_inverseN` rrules call one hand-written `_tensor_inverse_pullback`, in `src/kernels/inverses/tensor_inverse_pullback.jl`, and `scripts/utilities/inverse_generator.jl` emits that call. The tests are in `test/kernels/kernel_pullbacks.jl`: `test_rrule` in Float64, Float32 against a Float64 finite-difference reference, zero seeds, and seeds of the other precision. +- **The kernel rrules accept a zero seed and a seed of the other float type.** Each remaining kernel rrule (`tensor_tensor_mul`, `tensor_transpose_tensor_mul`, `tensor_mat_skew_sym_assign`, `tensor_inverse2` to `tensor_inverse5`, and `_mul(::Adjoint, x)`) answers a zero seed (`ZeroTangent`, `NoTangent` or a thunk of one) with `NoTangent()` or `ZeroTangent()`, where it raised a `MethodError`. It converts a seed of the other float type to the element type of the primal, where a Float64 seed on a Float32 primal raised a `MethodError`. `init_output` already did this. The style is the same in every rule: eager, with no `@thunk`. The four generated `tensor_inverseN` rrules call one hand-written `_tensor_inverse_pullback`, in `src/kernels/inverses/tensor_inverse_pullback.jl`, and `scripts/utilities/inverse_generator.jl` emits that call. The tests are in `test/kernels/kernel_pullbacks.jl`: `test_rrule` in Float64, Float32 against a Float64 finite-difference reference, zero seeds, and seeds of the other precision. - **A network with a Stiefel leaf trains on data of the other precision.** Training on data whose element type differs from the network's raised `AssertionError: eltype(A) == eltype(B)` in `mat_tensor_mul`. `_mul` checks no element type, because GEMM promotes. This closes known issue K7, which leaves `KNOWN_ISSUES.md`. The test is `test/integration/mixed_precision_training.jl`. @@ -98,6 +98,47 @@ breaking release). ### Changed +- **`VolumePreservingAttention` computes the Cayley transform in one kernel, `tensor_cayley`, at every sequence length and on a device array; the batch inverse that served the other lengths is deleted, and the closed forms for N = 2 to 5 stay, with N = 4 and 5 waiting for the user's ruling.** + + `src/kernels/cayley.jl` is new. `tensor_cayley(C::AbstractArray{T, 3})` is internal and not exported. It computes `(I + Cₖ)⁻¹(I − Cₖ)` for every slice `k` by row-pivoted Gauss–Jordan elimination on the pair `(I + Cₖ, I − Cₖ)`, one work item per slice, with the loop bound read from `size(C, 1)` at run time. Its `rrule` computes `C̄ = −(I + Λ)ᵀ Λ̄ (I + Λ)ᵀ / 2` from `Λ` alone: no second solve, no identity tensor, and no `Zygote.pullback` of `inv` inside a kernel. The activation of `VolumePreservingAttention` for every sequence length but 2 to 5 runs this kernel, so those lengths run on a device; the path it replaces, `cpu_inverse`, called `inv` and `Zygote.pullback(inv, ·)` on one slice inside a kernel, which this change removes. + + `cpu_inverse`, its `rrule` and `cpu_tensor_cayley` are deleted, because nothing calls them after this change. The identity tensor that the closed forms use moved from `src/kernels/exponentials/tensor_exponential.jl` to `src/kernels/identity_tensor.jl`, which holds `init_output`, `assign_ones!` and the `rrule` of `init_output`, so the file name matches what it holds, and known issue C18 leaves `KNOWN_ISSUES.md`. + + `scripts/utilities/cayley_timings.jl` is new. It times `tensor_cayley` on a batch of 1000 skew slices at scale 1 against the closed form `tensor_cayleyN` of the same `N`, for `N` = 2 to 8 and both element types. Each variant, `N` and element type runs in its own fresh process, with BLAS pinned to one thread and one untimed warm-up call; a time is the median of five calls, forward and backward. The script prints the kernel's two times and the ratio of the kernel to the closed form. For `N` = 6 to 8 there is no closed form in the tree, so it prints the kernel's times alone. + + Check 1, the medians the decision rests on, was run at commit `76cee621`, the commit of this branch that still held every closed form, including the batch inverse of `cpu_tensor_cayley`, which the N = 6 to 8 ratios are against. The shipped script cannot reproduce that column, because `cpu_tensor_cayley` is deleted. Machine: Apple M-series, Julia 1.13.2, shared with other work, at a load average between 49 and 88 at the start and the end of the run; BLAS on one thread; batch of 1000 skew slices at scale 1. The two µs columns are the kernel's own forward and backward times, and each ratio is the kernel over the closed form. + + | N | T | forward kernel µs | backward kernel µs | ratio forward | ratio backward | + |:--|:--|--:|--:|--:|--:| + | 2 | Float32 | 45.3 | 43.0 | 0.47 | 0.63 | + | 3 | Float32 | 74.7 | 72.3 | 1.71 | 0.58 | + | 4 | Float32 | 123.8 | 135.1 | 0.59 | 0.66 | + | 5 | Float32 | 184.5 | 224.3 | 1.43 | 0.66 | + | 6 | Float32 | 606.8 | 337.0 | 0.64 | 0.26 | + | 7 | Float32 | 372.0 | 485.5 | 0.33 | 0.30 | + | 8 | Float32 | 521.4 | 656.9 | 0.38 | 0.33 | + | 2 | Float64 | 48.8 | 51.0 | 1.46 | 0.65 | + | 3 | Float64 | 75.4 | 100.7 | 1.23 | 0.69 | + | 4 | Float64 | 124.4 | 168.6 | 1.13 | 0.62 | + | 5 | Float64 | 185.2 | 245.0 | 1.27 | 0.60 | + | 6 | Float64 | 289.8 | 384.5 | 0.16 | 0.23 | + | 7 | Float64 | 416.1 | 533.4 | 0.30 | 0.28 | + | 8 | Float64 | 565.8 | 756.6 | 0.34 | 0.33 | + + Run again on the shipped script at a lower load average (14.8 to 25.0), with the closed forms it + still holds, the decision is the same. The kernel's own times agree within 3 % of the table above, + except one N = 6 Float32 forward reading, 606.8 → 263.1. The closed form's own time is what moves, + by up to 3.3× (96 µs to 28.8 µs at N = 2, Float32, forward), so one median at these sizes is not a + figure of the code alone. The rows the decision rests on are above 1 in both runs: N = 2, Float32 + forward 0.47 then 1.51 and Float64 forward 1.46 then 1.33; N = 3, Float32 forward 1.71 then 1.79 + and Float64 forward 1.23 then 1.68; N = 4, Float32 forward 0.59 then 1.70 and Float64 forward 1.13 + in both; N = 5, Float32 forward 1.43 then 1.42 and Float64 forward 1.27 then 1.23. Every backward + ratio is below 1 in both runs. + + The orthogonality check in `test/kernels/cayley.jl` compares the kernel's `ΛᵀΛ − I` with the pivoted reference `(I + Cₖ) \ (I − Cₖ)` over `N` = 2 to 8, both element types and two scales, on three input sets. The worst ratio of the kernel's error to the reference's, against a limit of 4, is 1.54 for a hand-built skew tensor `R − Rᵀ`, 2.45 for the layer's `:arbitrary` path and 1.30 for its `:skew_sym` path. On the two exactly skew sets the kernel's own error is within the bound `10·n·eps(T)·(1 + ‖Cₖ‖₂²)`, its worst at 0.100 of the bound; on `:skew_sym`, where `C` is skew only to rounding, the pivoted reference itself exceeds that bound by a factor of 3.2, so that set carries the ratio clause alone. + + The closed forms `tensor_cayley2` to `tensor_cayley5` and their dispatch stay. A closed form stays while the kernel is slower than it at a median ratio above 1, and N = 6 to 8 are replaced whatever their ratio. N = 2 stays on its Float64 forward ratio of 1.46. N = 3 stays on its Float32 forward ratio of 1.71 and its Float64 forward ratio of 1.23. N = 4 stays until the user rules, on its Float64 forward ratio of 1.13, and N = 5 stays until the user rules, on its Float64 forward ratio of 1.27 and its Float32 forward ratio of 1.43. For N = 4 and 5 the question the numbers leave open is the user's to decide: keep the closed forms, or take the line savings of replacing them with the kernel. N = 6 to 8 are replaced: the kernel's ratios there run from 0.16 to 0.64. + - `Pkg.test()` no longer runs the doctests. `test/quality/doctests.jl` is now the group `doctests`, which an empty `ARGS` does not run; `Pkg.test(test_args = ["doctests"])` runs it. In CI the Doctests job stays their runner, so the test matrix no longer runs them a second time. - **CI uploads coverage from the `Julia 1 - ubuntu-latest` job instead of `Julia min`, and a test job saves the Julia cache only when it succeeds.** - **The package needs Julia 1.12 or later** (`[compat] julia = "1.12"`, was `"1.11"`). On Julia 1.11.9, the gradient of an `LNNLoss` with respect to the network parameters stops with `StackOverflowError` in type inference ("Internal error: during type inference of … Encountered stack overflow") on macOS and Windows; the `Julia min` job on Linux passes. The overflow comes from SymbolicUtils 4.49.0, whose `rrule` for `SymbolicUtils.Code.create_array` returns an exactly typed tuple of tangents, one per array element; Zygote's reverse pass through the generated Hessian of `lagrangian_acceleration` is then inferred far deeper. With the 4.48.0 rule restored, or in a task with a 64 MB stack, the same gradient passes on 1.11.9; Julia 1.12.7 and 1.13.1 pass it on the root task. This closes #327: the `Julia min` jobs on macOS and Windows had failed in the *Lagrangian Neural Network* testset since 2026-10-01, and every testset after it was never run there. @@ -108,7 +149,7 @@ breaking release). - **The optimizer reads `show_progress` as a `Bool`** instead of comparing it with `== true`. A non-`Bool` value now raises a `TypeError`: before, `nothing`, `0` or a string hid the progress bar with no error, and `1` or `1.0` showed it. Every call in the tests and the manual passes `false`, and every call in `scripts/` passes `true`. - **fatou reports no finding in `src/` or `test/`.** Only tests and one lint suppression changed; no behaviour of the package changed. `test/layers/multi_head_attention_stiefel_setup.jl` defined `check_setup(ps::NetworkParameters)` twice, and the later method silently replaced the `mapparameters` one, which never ran. The later method is deleted, and the test now asserts that the matrix method of `check_setup` runs once per leaf of the parameters, counted with `foldparameters`. The unused destructured values `Σ`, `Vt` and `Ũ₂` in `test/optimizers/psd_optim.jl` and `test/optimizers/svd_optim.jl` are now `_`; those files check `Ũ₁` and `Ũ₃` and never checked `Ũ₂`. The unreachable `new{…}(…)` after the `error` in the `GeneralizedHamiltonianArchitecture` constructor stays, because it records the field order for the future constructor, under `# fatou-ignore unreachable-code`. - **Three test files left the top level of `test/`**, because a test file sits there only where it mirrors `src/.jl`. `test/norm_eltype.jl` tests `_norm` of `src/utils.jl` and is now `test/utils.jl`; `test/map_to_cpu_tests.jl` tests `src/map_to_cpu.jl` and is now `test/map_to_cpu.jl`; `test/documented_defaults.jl` tests the keyword defaults of `src/layers/symplectic_attention.jl`, `src/architectures/symplectic_transformer.jl` and `src/architectures/standard_transformer_integrator.jl`, whose deepest common directory is `src/` itself, and is now `test/integration/documented_defaults.jl`. Only the paths in `test/runtests.jl` changed; the labels, groups and order are the same. -- **The matrix and tensor products are one internal product, `_mul`.** `src/kernels/mul.jl` defines `_mul(W, x)` as `W * x` for a vector or matrix `x`, so a structured weight is multiplied by its own `*`, and as one GEMM over `reshape(x, size(x, 1), :)` for a 3-tensor `x`; `_mul(x3, W)` as one GEMM between two `permutedims` of `x3`; and `_mul(x3, y3) = tensor_tensor_mul`. It replaces `mat_tensor_mul` and `mat_tensor_mul!`, `tensor_mat_mul` and `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. `custom_vec_mul` is `.*` now, and `vec_tensor_mul` and `tensor_scalar_product` are gone; a broadcast, and Zygote's rule for it, replaces them. A `SymmetricMatrix`, `SkewSymMatrix`, `StrictlyLowerTriangular` or `StrictlyUpperTriangular` weight is made dense once per call, as `W * one(W)` on its own backend, in a 3-tensor left product and in every right product; a vector or matrix `x` in a left product takes the weight's own `*` instead. A structured matrix as the right argument of a left product is made dense the same way. Its gradient reaches the storage through GeometricOptimizers' `ProjectTo` and NeuralNetworkParameters' `storage_gradient`. A short rrule on `_mul(::Adjoint, x)` keeps the cotangent of a `Matrix` parameter that is used as `W'` a `Matrix`, at every rank of `x` and for a Stiefel leaf. `ManifoldLayer`, `ResNetLayer` and `VolumePreservingFeedForwardLayer` each have one `x::AbstractArray` method in place of a matrix method and a tensor method. `tensor_transpose`, with its kernel and rrule, is gone: callers write `permutedims(x, (2, 1, 3))`. `cpu_tensor_cayley` loses its shadowed duplicate in `cpu_inverse.jl`, and the four kernel files that survive lose the file-top strings that described the removed kernels. No exported name changes, so this is not a breaking change. +- **The matrix and tensor products are one internal product, `_mul`.** `src/kernels/mul.jl` defines `_mul(W, x)` as `W * x` for a vector or matrix `x`, so a structured weight is multiplied by its own `*`, and as one GEMM over `reshape(x, size(x, 1), :)` for a 3-tensor `x`; `_mul(x3, W)` as one GEMM between two `permutedims` of `x3`; and `_mul(x3, y3) = tensor_tensor_mul`. It replaces `mat_tensor_mul` and `mat_tensor_mul!`, `tensor_mat_mul` and `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. `custom_vec_mul` is `.*` now, and `vec_tensor_mul` and `tensor_scalar_product` are gone; a broadcast, and Zygote's rule for it, replaces them. A `SymmetricMatrix`, `SkewSymMatrix`, `StrictlyLowerTriangular` or `StrictlyUpperTriangular` weight is made dense once per call, as `W * one(W)` on its own backend, in a 3-tensor left product and in every right product; a vector or matrix `x` in a left product takes the weight's own `*` instead. A structured matrix as the right argument of a left product is made dense the same way. Its gradient reaches the storage through GeometricOptimizers' `ProjectTo` and NeuralNetworkParameters' `storage_gradient`. A short rrule on `_mul(::Adjoint, x)` keeps the cotangent of a `Matrix` parameter that is used as `W'` a `Matrix`, at every rank of `x` and for a Stiefel leaf. `ManifoldLayer`, `ResNetLayer` and `VolumePreservingFeedForwardLayer` each have one `x::AbstractArray` method in place of a matrix method and a tensor method. `tensor_transpose`, with its kernel and rrule, is gone: callers write `permutedims(x, (2, 1, 3))`. The four kernel files that survive lose the file-top strings that described the removed kernels. No exported name changes, so this is not a breaking change. Measured with `scripts/utilities/product_timings.jl` at 10×10×1000, in µs, before → after. The script's `x` is a 3-tensor, so this table does not cover a vector or matrix `x`. The before column is the same script run at `2a63f26a` with the old names, `mat_tensor_mul` for the left product and `tensor_mat_mul` for the right, with base and branch processes alternated. BLAS ran on one thread. Each variant ran in five fresh processes, each doing one warm-up and taking the minimum of 100 calls, and the table gives the median of the five. The backward column is `Zygote.gradient` through a `NetworkParameters`. The machine was an Apple M-series with Julia 1.13.1, GeometricOptimizers 0.9.0, NeuralNetworkParameters 0.5.2 and Zygote 0.7.13, shared with other work, at a load average between 8 and 37. diff --git a/KNOWN_ISSUES.md b/KNOWN_ISSUES.md index bc059850d..9d53fad06 100644 --- a/KNOWN_ISSUES.md +++ b/KNOWN_ISSUES.md @@ -89,16 +89,6 @@ What is known to be wrong in GeometricMachineLearning and is not fixed. deliberately, and the resolution is upstream declaring them `public`. Each check's reason is written beside it in the test file. -### C18 · `src/kernels/exponentials/tensor_exponential.jl` no longer defines `tensor_exponential`; it holds the identity tensor and its `rrule`, used by `tensor_cayley.jl` and `cpu_inverse.jl`. - -- **location:** `src/kernels/exponentials/tensor_exponential.jl` -- **kind:** defect -- **found:** 2026-09-19 -- **evidence:** - The file name, and the `exponentials/` directory around it, are wider than what they hold. A - comment in the file itself cites "issue C18" by number, so this entry exists to keep that reference - valid. Renaming both the file and the directory is the fix. - ### C19 · `_make_block_for_initialization` has two methods - **location:** `src/architectures/symplectic_transformer.jl` diff --git a/scripts/utilities/cayley_timings.jl b/scripts/utilities/cayley_timings.jl index 86b236d7a..a9b5e5be4 100644 --- a/scripts/utilities/cayley_timings.jl +++ b/scripts/utilities/cayley_timings.jl @@ -1,28 +1,30 @@ # The forward and the backward time of the Cayley transform of a batch of 1000 skew slices at scale -# 1, `tensor_cayley`, against the closed forms it would replace: `tensor_cayleyN` for N = 2 to 5 and -# `cpu_tensor_cayley` for N = 6 to 8, each where the tree still defines it. BLAS is pinned to one -# thread, and every variant, every N and every element type runs in a fresh process; a process makes -# one warm-up call it does not time, then takes the median of five timed calls. The parent prints the -# ratio of the kernel to the closed form, forward and backward. +# 1, `tensor_cayley`, against the closed form the tree defines for that slice size: `tensor_cayleyN` +# for N = 2 to 5. A size whose closed form the tree does not define, N = 6 to 8, prints the kernel's +# own time and no ratio. BLAS is pinned to one thread, and every variant, every N and every element +# type runs in a fresh process; a process makes one warm-up call it does not time, then takes the +# median of five timed calls. The parent prints the ratio of the kernel to the closed form, forward +# and backward. # -# The closed forms that the kernel replaces are deleted by the outcome this table decides, so the -# figures come from a branch commit that holds both the kernel and every closed form, and the -# `CHANGELOG.md` entry of that outcome cites the hash of that commit. After the deletions the script -# times only the variants still defined (`isdefined`). +# The table that decided which closed forms stay carried a column for N = 6 to 8 as well, against the +# batch inverse of a long slice, which the tree no longer defines. That column came from a commit +# that still held it, and the `CHANGELOG.md` entry of the change names the commit. # -# Run from the repository root, with the scripts environment; one process of one variant is the -# second form, which the first starts: -# julia --startup-file=no --project=scripts scripts/utilities/cayley_timings.jl -# julia --startup-file=no --project=scripts scripts/utilities/cayley_timings.jl tensor_cayley 4 Float32 +# Run from the repository root, with the package's own environment, which carries `Zygote` and the +# standard libraries this script uses; `--project=scripts` resolves them as well, and CUDA with them: +# julia --startup-file=no --project=. scripts/utilities/cayley_timings.jl +# julia --startup-file=no --project=. scripts/utilities/cayley_timings.jl tensor_cayley 4 Float32 # -# Metal has no reference for N = 6 to 8, because `cpu_tensor_cayley` calls `inv` inside a kernel, and -# Metal is not a dependency of `scripts/Project.toml`. The Metal column runs `measure` through -# Kaimon, in a scratch environment that develops the worktree and adds Metal, as +# Metal has no closed form to compare against, because the symbolic inverses cover 2 to 5 only, and +# `Metal` is not a dependency of the package. Its column runs `measure` through Kaimon, in a scratch +# environment that develops this worktree and adds Metal. This file runs its table only when it is +# not included, so a session reaches `measure` by including it and calling it: # +# julia> include("scripts/utilities/cayley_timings.jl") # julia> measure("tensor_cayley", 6, Float32, Metal.@sync) # -# because a Metal launch returns before the kernel has run. That column prints the kernel's own time -# and no ratio. +# `Metal.@sync` is what makes the figure a time rather than a launch: without it the call returns +# before the kernel has run. using GeometricMachineLearning using LinearAlgebra: BLAS @@ -64,19 +66,15 @@ function measure(variant, n, ::Type{T}, sync = identity) where {T} forward, backward end -# The closed form of the transform of an `n × n` slice, where the tree still defines it. -function reference(n) - n <= 5 ? "tensor_cayley$n" : "cpu_tensor_cayley" -end +# The closed form of the transform of an `n × n` slice, where the tree defines one. +reference(n) = n <= 5 ? "tensor_cayley$n" : nothing function one_process(variant, n, T) cmd = `$(Base.julia_cmd()) --startup-file=no --project=$(Base.active_project()) $(@__FILE__) $variant $n $T` parse.(Float64, split(readchomp(cmd))) end -if length(ARGS) == 3 - println(join(measure(ARGS[1], parse(Int, ARGS[2]), getfield(Base, Symbol(ARGS[3]))), " ")) -else +function table() @printf("%-17s %-2s %-8s %11s %11s %s\n", "variant", "N", "T", "forward µs", "backward µs", "ratio kernel/closed form") @@ -84,7 +82,7 @@ else kernel = one_process(KERNEL, n, T) name = reference(n) - if isdefined(GeometricMachineLearning, Symbol(name)) + if name !== nothing && isdefined(GeometricMachineLearning, Symbol(name)) closed = one_process(name, n, T) @printf("%-17s %-2d %-8s %11.1f %11.1f forward %.2f, backward %.2f\n", KERNEL, n, T, kernel[1], kernel[2], kernel[1] / closed[1], @@ -95,3 +93,9 @@ else end end end + +if length(ARGS) == 3 + println(join(measure(ARGS[1], parse(Int, ARGS[2]), getfield(Base, Symbol(ARGS[3]))), " ")) +elseif !isinteractive() + table() +end diff --git a/scripts/utilities/inverse_generator.jl b/scripts/utilities/inverse_generator.jl index 145fe96bf..809d30b32 100644 --- a/scripts/utilities/inverse_generator.jl +++ b/scripts/utilities/inverse_generator.jl @@ -81,7 +81,7 @@ function wrapper_source(n) end # The four sizes the package ships. `VolumePreservingAttention` dispatches on the sequence length -# and falls back to `cpu_tensor_cayley` for every other one. +# and falls back to `tensor_cayley` for every other one. for n in 2:5 println("$(n)x$(n)...") write("inverse_$(n)x$(n).jl", kernel_source(n) * "\n\n" * wrapper_source(n)) diff --git a/src/GeometricMachineLearning.jl b/src/GeometricMachineLearning.jl index d6fb9af43..379913c8d 100644 --- a/src/GeometricMachineLearning.jl +++ b/src/GeometricMachineLearning.jl @@ -128,8 +128,7 @@ include("kernels/tensor_transpose_tensor_mul.jl") include("kernels/tensor_tensor_transpose_mul.jl") include("kernels/mul.jl") include("kernels/cayley.jl") -include("kernels/exponentials/tensor_exponential.jl") -include("kernels/inverses/cpu_inverse.jl") +include("kernels/identity_tensor.jl") include("kernels/inverses/tensor_inverse_pullback.jl") include("kernels/inverses/inverse_2x2.jl") include("kernels/inverses/inverse_3x3.jl") diff --git a/src/kernels/exponentials/tensor_exponential.jl b/src/kernels/identity_tensor.jl similarity index 80% rename from src/kernels/exponentials/tensor_exponential.jl rename to src/kernels/identity_tensor.jl index 663dbfc92..b53fcdbf2 100644 --- a/src/kernels/exponentials/tensor_exponential.jl +++ b/src/kernels/identity_tensor.jl @@ -1,5 +1,5 @@ -# The identity tensor and its `rrule`, used by `tensor_cayley.jl` and `cpu_inverse.jl`. The file -# name is wider than what it holds; renaming it, and `exponentials/` with it, is issue C18. +# The identity tensor and its `rrule`, which the closed forms of the Cayley transform of a batch of +# small slices, `tensor_cayley2` to `tensor_cayley5`, use to form `I ± C`. function init_output(B::AbstractArray{T, 3}) where {T} output = zero(B) diff --git a/src/kernels/inverses/cpu_inverse.jl b/src/kernels/inverses/cpu_inverse.jl deleted file mode 100644 index f21cdc478..000000000 --- a/src/kernels/inverses/cpu_inverse.jl +++ /dev/null @@ -1,48 +0,0 @@ -@kernel function cpu_inverse_kernel!(B, A) - k = @index(Global) - @views A_temp = A[:, :, k] - @views B_temp = B[:, :, k] - - B_temp .= inv(A_temp) - - nothing -end - -function cpu_inverse(A::AbstractArray) - B = zero(A) - backend = networkbackend(A) - - cpu_inverse! = cpu_inverse_kernel!(backend) - cpu_inverse!(B, A, ndrange = size(A, 3)) - - B -end - -@kernel function cpu_inverse_pullback_kernel!(dA, A, dB) - k = @index(Global) - @views A_temp = A[:, :, k] - @views dA_temp = dA[:, :, k] - @views dB_temp = dB[:, :, k] - - copy!(dA_temp, Zygote.pullback(inv, A_temp)[2](dB_temp)[1]) - - nothing -end - -function ChainRulesCore.rrule(::typeof(cpu_inverse), A::AbstractArray) - B = cpu_inverse(A) - - function cpu_inverse_pullback(dB) - dB = _seed(eltype(A), dB) - dB isa AbstractZero && return NoTangent(), ZeroTangent() - dA = zero(dB) - backend = networkbackend(dB) - - cpu_inverse_pullback! = cpu_inverse_pullback_kernel!(backend) - cpu_inverse_pullback!(dA, A, dB, ndrange = size(dB, 3)) - - return NoTangent(), dA - end - - return B, cpu_inverse_pullback -end diff --git a/src/kernels/inverses/tensor_cayley.jl b/src/kernels/inverses/tensor_cayley.jl index d234d8d29..473641534 100644 --- a/src/kernels/inverses/tensor_cayley.jl +++ b/src/kernels/inverses/tensor_cayley.jl @@ -1,8 +1,3 @@ -function cpu_tensor_cayley(A::AbstractArray{T, 3}) where {T} - one_tensor = init_output(A) - tensor_tensor_mul(one_tensor - A, cpu_inverse(one_tensor + A)) -end - function tensor_cayley2(A::AbstractArray{T, 3}) where {T} one_tensor = init_output(A) tensor_tensor_mul(one_tensor - A, tensor_inverse2(one_tensor + A)) diff --git a/src/layers/volume_preserving_attention.jl b/src/layers/volume_preserving_attention.jl index e862928f4..80e074a92 100644 --- a/src/layers/volume_preserving_attention.jl +++ b/src/layers/volume_preserving_attention.jl @@ -3,9 +3,9 @@ Make an instance of `VolumePreservingAttention` for a specific dimension and sequence length. -The sequence length is `0` by default. - -Setting `seq_length` to `0` for all sequence lengths but does not apply the fast Cayley activation. +The sequence length is `0` by default, and a layer built with `0` applies to data of any sequence +length: the Cayley transform of the data is one kernel, [`tensor_cayley`](@ref), rather than a +symbolic inverse, and every sequence length runs on the CPU and on a device. # Arguments @@ -20,10 +20,12 @@ Applying a layer of type `VolumePreservingAttention` does the following: # Implementation -The fast activation is only implemented for sequence lengths of 2, 3, 4 and 5. -Other sequence lengths only work on CPU (for now). +The activation is one kernel, [`tensor_cayley`](@ref), which solves ``(\mathbb{I} + C)\Lambda = +\mathbb{I} - C`` by row-pivoted Gauss–Jordan elimination, one work item per slice: it runs at any +sequence length, on the CPU and on a device, and in the element type of the data. -The fast Cayley activation is using inverses that have been computed symbolically. +The sequence lengths 2, 3, 4 and 5 use the closed forms `tensor_cayley2` to `tensor_cayley5` +instead, whose inverses have been computed symbolically. """ struct VolumePreservingAttention{M, N, ScalarProductType, SL} <: AbstractExplicitLayer{M, N} function VolumePreservingAttention(dim::Int, seq_length::Int = 0; skew_sym::Bool = false) diff --git a/test/kernels/inverses/test_cayley_transforms.jl b/test/kernels/inverses/test_cayley_transforms.jl index 9c4b085a1..bb80dd15a 100644 --- a/test/kernels/inverses/test_cayley_transforms.jl +++ b/test/kernels/inverses/test_cayley_transforms.jl @@ -1,4 +1,4 @@ -using GeometricMachineLearning: tensor_cayley4, tensor_cayley3, cpu_tensor_cayley +using GeometricMachineLearning: tensor_cayley4, tensor_cayley3 using Test import Random @@ -23,8 +23,6 @@ function test_cayley_transforms(third_dim::Int = 10) A₃ = A₃ - permutedims(A₃, (2, 1, 3)) test_orthonormal(tensor_cayley4(A₄)) test_orthonormal(tensor_cayley3(A₃)) - test_orthonormal(cpu_tensor_cayley(A₄)) - test_orthonormal(cpu_tensor_cayley(A₃)) end test_cayley_transforms() diff --git a/test/kernels/kernel_pullbacks.jl b/test/kernels/kernel_pullbacks.jl index 9351d1b38..98d62b0f3 100644 --- a/test/kernels/kernel_pullbacks.jl +++ b/test/kernels/kernel_pullbacks.jl @@ -1,5 +1,5 @@ using GeometricMachineLearning: tensor_tensor_mul, tensor_transpose_tensor_mul, - tensor_mat_skew_sym_assign, cpu_inverse, tensor_inverse2, + tensor_mat_skew_sym_assign, tensor_inverse2, tensor_inverse3, tensor_inverse4, tensor_inverse5, init_output, tensor_cayley, _mul @@ -33,7 +33,6 @@ function main(first_dim, second_dim, third_dim, third_tensor_dim) rand(first_dim, first_dim)) test_rrule(_mul, rand(first_dim, second_dim)', rand(first_dim, third_dim, third_tensor_dim)) test_rrule(_mul, rand(first_dim, second_dim)', rand(first_dim, third_dim)) - test_rrule(cpu_inverse, invertible(Float64, first_dim, third_tensor_dim)) test_rrule(init_output, rand(first_dim, first_dim, third_tensor_dim)) test_rrule(tensor_cayley, skew(Float64, first_dim, third_tensor_dim)) end @@ -62,7 +61,6 @@ function rules(T) (tensor_transpose_tensor_mul, (rand(T, 4, 3, 2), rand(T, 4, 5, 2))), (tensor_mat_skew_sym_assign, (rand(T, 3, 4, 2), rand(T, 3, 3))), (_mul, (rand(T, 4, 3)', rand(T, 4, 5, 2))), - (cpu_inverse, (invertible(T, 3, 2),)), (tensor_inverse2, (invertible(T, 2, 2),)), (tensor_inverse3, (invertible(T, 3, 2),)), (tensor_inverse4, (invertible(T, 4, 2),)), diff --git a/test/kernels/tensor_inverse.jl b/test/kernels/tensor_inverse.jl index 9a556a1d6..82a82ec34 100644 --- a/test/kernels/tensor_inverse.jl +++ b/test/kernels/tensor_inverse.jl @@ -1,5 +1,5 @@ using GeometricMachineLearning: tensor_inverse2, tensor_inverse3, tensor_inverse4, - tensor_inverse5, cpu_inverse + tensor_inverse5 using Test import Zygote import Random @@ -121,21 +121,3 @@ function test22_inverse_pullback(k::Int = 10) end test22_inverse_pullback() - -function test_cpu_inverse_pullback(k::Int = 10) - A = rand(3, 3, k) - - pullback_total = Zygote.pullback(cpu_inverse, A) - - out_diff = rand(3, 3, k) - in_diff = pullback_total[2](out_diff)[1] - - for i in 1:k - pullback_i = Zygote.pullback(inv, A[:, :, i]) - - @test pullback_total[1][:, :, i] ≈ pullback_i[1] - @test in_diff[:, :, i] ≈ pullback_i[2](out_diff[:, :, i])[1] - end -end - -test_cpu_inverse_pullback() diff --git a/test/quality/jet.jl b/test/quality/jet.jl index 19713b244..bbc9844ea 100644 --- a/test/quality/jet.jl +++ b/test/quality/jet.jl @@ -58,7 +58,7 @@ const A3{T} = Array{T, 3} (2, 3, 2), zeros(T, 2, 3, 2), zeros(T, 2, 4, 2), zeros(T, 3, 4, 2))) end - # src/kernels/exponentials/tensor_exponential.jl + # src/kernels/identity_tensor.jl for T in (Float16, Float32, Float64) @test isempty(JET.get_reports(JET.report_opt(GML.assign_ones!, (A3{T},); target_modules = TARGET))) @@ -91,16 +91,7 @@ const A3{T} = Array{T, 3} @test isempty(kernel_body_reports(GML.da_kernel!, (3, 3, 2), zeros(3, 3, 2), Z, M, zeros(4, 4, 2), 3, 4)) - # src/kernels/inverses/: `cpu_inverse` and its pullback in test/kernels/tensor_inverse.jl - @test isempty(JET.get_reports(JET.report_opt(GML.cpu_inverse, (A3{Float64},); - target_modules = TARGET))) - @test isempty(kernel_body_reports(GML.cpu_inverse_kernel!, 2, zero(B), B)) - Id = zeros(3, 3, 2) - foreach(i -> Id[i, i, :] .= 1, 1:3) - @test isempty(JET.get_reports(JET.report_opt(rrule(GML.cpu_inverse, Id)[2], - (A3{Float64},); target_modules = TARGET))) - @test isempty(kernel_body_reports(GML.cpu_inverse_pullback_kernel!, 2, zero(B), B, - dC)) + # src/kernels/inverses/: the symbolic inverses in test/kernels/tensor_inverse.jl for T in (Float16, Float32, Float64) for (f, k, n) in ((GML.tensor_inverse2!, GML.inv22_kernel!, 2), (GML.tensor_inverse3!, GML.inv33_kernel!, 3), From c7a0266ceb4f16aa224ef1727d382f0f0f876c7b Mon Sep 17 00:00:00 2001 From: Michael Kraus Date: Fri, 16 Oct 2026 23:22:29 +0200 Subject: [PATCH 03/10] GML part M4: drop the jet.jl locals the deleted rows held `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. --- test/quality/jet.jl | 1 - 1 file changed, 1 deletion(-) diff --git a/test/quality/jet.jl b/test/quality/jet.jl index bbc9844ea..bfa514642 100644 --- a/test/quality/jet.jl +++ b/test/quality/jet.jl @@ -81,7 +81,6 @@ const A3{T} = Array{T, 3} # src/kernels/kernel_ad_routines/: the pullback that launches kernels, at the `Float64` # types of `test_rrule` in test/kernels/kernel_pullbacks.jl - B, dC = zeros(3, 3, 2), zeros(3, 3, 2) Z, M = zeros(3, 4, 2), zeros(3, 3) @test isempty(JET.get_reports(JET.report_opt( rrule(GML.tensor_mat_skew_sym_assign, Z, M)[2], (A3{Float64},); From 632ae810b87c407cd23b247f62d0c1e1259277a7 Mon Sep 17 00:00:00 2001 From: Michael Kraus Date: Fri, 16 Oct 2026 23:25:21 +0200 Subject: [PATCH 04/10] GML part M4: the check 1 evidence, the reproducibility run and the Metal column MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- CHANGELOG.md | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 0d1ecbcce..270ef45ff 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -104,6 +104,8 @@ breaking release). `cpu_inverse`, its `rrule` and `cpu_tensor_cayley` are deleted, because nothing calls them after this change. The identity tensor that the closed forms use moved from `src/kernels/exponentials/tensor_exponential.jl` to `src/kernels/identity_tensor.jl`, which holds `init_output`, `assign_ones!` and the `rrule` of `init_output`, so the file name matches what it holds, and known issue C18 leaves `KNOWN_ISSUES.md`. + `test/kernels/cayley.jl` is new: the exact skewness of the layer's `:arbitrary` path, the kernel's orthogonality error against the pivoted reference and against the bound `10·n·eps(T)·(1 + ‖Cₖ‖₂²)`, its value against a `BigFloat` solve, and its pullback against a central difference. `test/layers/volume_preserving_attention.jl` is new: 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` at sequence lengths 3 and 7. `test/devices/jlarray.jl` and `test/devices/metal.jl` gain the kernel's device rows, the Metal one a `Float32` testset that also runs the layer's forward pass at N = 2 to 8, and `test/quality/jet.jl` reaches the kernel, its kernel body and its pullback in `Float16`, `Float32` and `Float64`. + `scripts/utilities/cayley_timings.jl` is new. It times `tensor_cayley` on a batch of 1000 skew slices at scale 1 against the closed form `tensor_cayleyN` of the same `N`, for `N` = 2 to 8 and both element types. Each variant, `N` and element type runs in its own fresh process, with BLAS pinned to one thread and one untimed warm-up call; a time is the median of five calls, forward and backward. The script prints the kernel's two times and the ratio of the kernel to the closed form. For `N` = 6 to 8 there is no closed form in the tree, so it prints the kernel's times alone. Check 1, the medians the decision rests on, was run at commit `76cee621`, the commit of this branch that still held every closed form, including the batch inverse of `cpu_tensor_cayley`, which the N = 6 to 8 ratios are against. The shipped script cannot reproduce that column, because `cpu_tensor_cayley` is deleted. Machine: Apple M-series, Julia 1.13.2, shared with other work, at a load average between 49 and 88 at the start and the end of the run; BLAS on one thread; batch of 1000 skew slices at scale 1. The two µs columns are the kernel's own forward and backward times, and each ratio is the kernel over the closed form. @@ -133,7 +135,15 @@ breaking release). forward 0.47 then 1.51 and Float64 forward 1.46 then 1.33; N = 3, Float32 forward 1.71 then 1.79 and Float64 forward 1.23 then 1.68; N = 4, Float32 forward 0.59 then 1.70 and Float64 forward 1.13 in both; N = 5, Float32 forward 1.43 then 1.42 and Float64 forward 1.27 then 1.23. Every backward - ratio is below 1 in both runs. + ratio on the CPU is below 1 in both runs. + + The Metal column, `Float32`, with every timed call wrapped in `Metal.@sync` and the same `measure`, + on an Apple M4 Max (`AGXG16CDevice`), through the test environment, which carries Metal where + `scripts/Project.toml` does not: the ratio of the kernel to the closed form is 0.61 forward and + 5.56 backward at N = 2, 0.41 and 0.71 at N = 3, 1.33 and 1.53 at N = 4, and 1.78 and 0.95 at + N = 5; the kernel's own forward time at N = 6, 7 and 8 is 2544.8, 1771.5 and 2297.3 µs, against + 292.3 at N = 2. The N = 2 backward ratio of 5.56 is the row check 1's rule names for the GPU: the + closed form for N = 2 stays on the CPU columns alone, and on this column too. The orthogonality check in `test/kernels/cayley.jl` compares the kernel's `ΛᵀΛ − I` with the pivoted reference `(I + Cₖ) \ (I − Cₖ)` over `N` = 2 to 8, both element types and two scales, on three input sets. The worst ratio of the kernel's error to the reference's, against a limit of 4, is 1.54 for a hand-built skew tensor `R − Rᵀ`, 2.45 for the layer's `:arbitrary` path and 1.30 for its `:skew_sym` path. On the two exactly skew sets the kernel's own error is within the bound `10·n·eps(T)·(1 + ‖Cₖ‖₂²)`, its worst at 0.100 of the bound; on `:skew_sym`, where `C` is skew only to rounding, the pivoted reference itself exceeds that bound by a factor of 3.2, so that set carries the ratio clause alone. From e873e20114257b047b6d888909fbf56970effc1d Mon Sep 17 00:00:00 2001 From: Michael Kraus Date: Fri, 16 Oct 2026 23:35:25 +0200 Subject: [PATCH 05/10] GML part M4: the ruling keeps the closed forms for 4 and 5 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. --- CHANGELOG.md | 14 +++++++------- scripts/utilities/cayley_timings.jl | 10 ++++++---- 2 files changed, 13 insertions(+), 11 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 270ef45ff..2acc0a1a1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -98,7 +98,7 @@ breaking release). ### Changed -- **`VolumePreservingAttention` computes the Cayley transform in one kernel, `tensor_cayley`, at every sequence length and on a device array; the batch inverse that served the other lengths is deleted, and the closed forms for N = 2 to 5 stay, with N = 4 and 5 waiting for the user's ruling.** +- **`VolumePreservingAttention` computes the Cayley transform in one kernel, `tensor_cayley`, at every sequence length and on a device array; the batch inverse that served the other lengths is deleted, and the closed forms for N = 2 to 5 stay.** `src/kernels/cayley.jl` is new. `tensor_cayley(C::AbstractArray{T, 3})` is internal and not exported. It computes `(I + Cₖ)⁻¹(I − Cₖ)` for every slice `k` by row-pivoted Gauss–Jordan elimination on the pair `(I + Cₖ, I − Cₖ)`, one work item per slice, with the loop bound read from `size(C, 1)` at run time. Its `rrule` computes `C̄ = −(I + Λ)ᵀ Λ̄ (I + Λ)ᵀ / 2` from `Λ` alone: no second solve, no identity tensor, and no `Zygote.pullback` of `inv` inside a kernel. The activation of `VolumePreservingAttention` for every sequence length but 2 to 5 runs this kernel, so those lengths run on a device; the path it replaces, `cpu_inverse`, called `inv` and `Zygote.pullback(inv, ·)` on one slice inside a kernel, which this change removes. @@ -131,11 +131,11 @@ breaking release). still holds, the decision is the same. The kernel's own times agree within 3 % of the table above, except one N = 6 Float32 forward reading, 606.8 → 263.1. The closed form's own time is what moves, by up to 3.3× (96 µs to 28.8 µs at N = 2, Float32, forward), so one median at these sizes is not a - figure of the code alone. The rows the decision rests on are above 1 in both runs: N = 2, Float32 - forward 0.47 then 1.51 and Float64 forward 1.46 then 1.33; N = 3, Float32 forward 1.71 then 1.79 - and Float64 forward 1.23 then 1.68; N = 4, Float32 forward 0.59 then 1.70 and Float64 forward 1.13 - in both; N = 5, Float32 forward 1.43 then 1.42 and Float64 forward 1.27 then 1.23. Every backward - ratio on the CPU is below 1 in both runs. + figure of the code alone. The forward ratios, first run then second: N = 2, Float32 0.47 then 1.51 + and Float64 1.46 then 1.33; N = 3, Float32 1.71 then 1.79 and Float64 1.23 then 1.68; N = 4, + Float32 0.59 then 1.70 and Float64 1.13 in both; N = 5, Float32 1.43 then 1.42 and Float64 1.27 + then 1.23. Every N from 2 to 5 has a forward ratio above 1 in both runs; every backward ratio on + the CPU is below 1 in both runs. The Metal column, `Float32`, with every timed call wrapped in `Metal.@sync` and the same `measure`, on an Apple M4 Max (`AGXG16CDevice`), through the test environment, which carries Metal where @@ -147,7 +147,7 @@ breaking release). The orthogonality check in `test/kernels/cayley.jl` compares the kernel's `ΛᵀΛ − I` with the pivoted reference `(I + Cₖ) \ (I − Cₖ)` over `N` = 2 to 8, both element types and two scales, on three input sets. The worst ratio of the kernel's error to the reference's, against a limit of 4, is 1.54 for a hand-built skew tensor `R − Rᵀ`, 2.45 for the layer's `:arbitrary` path and 1.30 for its `:skew_sym` path. On the two exactly skew sets the kernel's own error is within the bound `10·n·eps(T)·(1 + ‖Cₖ‖₂²)`, its worst at 0.100 of the bound; on `:skew_sym`, where `C` is skew only to rounding, the pivoted reference itself exceeds that bound by a factor of 3.2, so that set carries the ratio clause alone. - The closed forms `tensor_cayley2` to `tensor_cayley5` and their dispatch stay. A closed form stays while the kernel is slower than it at a median ratio above 1, and N = 6 to 8 are replaced whatever their ratio. N = 2 stays on its Float64 forward ratio of 1.46. N = 3 stays on its Float32 forward ratio of 1.71 and its Float64 forward ratio of 1.23. N = 4 stays until the user rules, on its Float64 forward ratio of 1.13, and N = 5 stays until the user rules, on its Float64 forward ratio of 1.27 and its Float32 forward ratio of 1.43. For N = 4 and 5 the question the numbers leave open is the user's to decide: keep the closed forms, or take the line savings of replacing them with the kernel. N = 6 to 8 are replaced: the kernel's ratios there run from 0.16 to 0.64. + The closed forms `tensor_cayley2` to `tensor_cayley5` and their dispatch stay. A closed form stays while the kernel is slower than it at a median ratio above 1, and N = 6 to 8 are replaced whatever their ratio. N = 2 stays on its Float64 forward ratio of 1.46. N = 3 stays on its Float32 forward ratio of 1.71 and its Float64 forward ratio of 1.23. N = 4 stays on its Float64 forward ratio of 1.13 in both runs and its Metal ratios of 1.33 forward and 1.53 backward. N = 5 stays on its Float32 forward ratio of 1.43 then 1.42, its Float64 forward ratio of 1.27 then 1.23 and its Metal forward ratio of 1.78. N = 4 and 5 are the user's call rather than the part's, because the kernel is slower forward there in every `Float64` run and on Metal; the user kept their closed forms, `inverse_4x4.jl` and `inverse_5x5.jl` with their 690 generated lines. N = 6 to 8 are replaced: the kernel's ratios there run from 0.16 to 0.64. - `Pkg.test()` no longer runs the doctests. `test/quality/doctests.jl` is now the group `doctests`, which an empty `ARGS` does not run; `Pkg.test(test_args = ["doctests"])` runs it. In CI the Doctests job stays their runner, so the test matrix no longer runs them a second time. - **CI uploads coverage from the `Julia 1 - ubuntu-latest` job instead of `Julia min`, and a test job saves the Julia cache only when it succeeds.** diff --git a/scripts/utilities/cayley_timings.jl b/scripts/utilities/cayley_timings.jl index a9b5e5be4..543e729cb 100644 --- a/scripts/utilities/cayley_timings.jl +++ b/scripts/utilities/cayley_timings.jl @@ -15,10 +15,12 @@ # julia --startup-file=no --project=. scripts/utilities/cayley_timings.jl # julia --startup-file=no --project=. scripts/utilities/cayley_timings.jl tensor_cayley 4 Float32 # -# Metal has no closed form to compare against, because the symbolic inverses cover 2 to 5 only, and -# `Metal` is not a dependency of the package. Its column runs `measure` through Kaimon, in a scratch -# environment that develops this worktree and adds Metal. This file runs its table only when it is -# not included, so a session reaches `measure` by including it and calling it: +# Metal has no reference for N = 6 to 8: the batch inverse that served those rows called `inv` inside +# a kernel, and the tree does not define it, so those rows print the kernel's time alone. N = 2 to 5 +# compare against the closed forms, which run on Metal as well; `Metal` is not a dependency of the +# package, so its column runs `measure` through Kaimon, in a scratch environment that develops this +# worktree and adds Metal. This file runs its table only when it is not included, so a session +# reaches `measure` by including it and calling it: # # julia> include("scripts/utilities/cayley_timings.jl") # julia> measure("tensor_cayley", 6, Float32, Metal.@sync) From 5ab69219a50d68fe9b82063aa77a661cbe2a4e3a Mon Sep 17 00:00:00 2001 From: Michael Kraus Date: Sat, 17 Oct 2026 01:54:24 +0200 Subject: [PATCH 06/10] GML part M4: the exact-skew row sees the layer's own operand, and the benchmark times a device MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- CHANGELOG.md | 30 ++++--- KNOWN_ISSUES.md | 25 ++++++ scripts/utilities/cayley_timings.jl | 52 +++++++----- src/kernels/cayley.jl | 12 ++- src/layers/volume_preserving_attention.jl | 20 +++-- test/devices/jlarray.jl | 26 ++++-- test/kernels/cayley.jl | 95 ++++++++++------------ test/layers/volume_preserving_attention.jl | 48 +++++++---- 8 files changed, 187 insertions(+), 121 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 2acc0a1a1..e3d369030 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -104,11 +104,11 @@ breaking release). `cpu_inverse`, its `rrule` and `cpu_tensor_cayley` are deleted, because nothing calls them after this change. The identity tensor that the closed forms use moved from `src/kernels/exponentials/tensor_exponential.jl` to `src/kernels/identity_tensor.jl`, which holds `init_output`, `assign_ones!` and the `rrule` of `init_output`, so the file name matches what it holds, and known issue C18 leaves `KNOWN_ISSUES.md`. - `test/kernels/cayley.jl` is new: the exact skewness of the layer's `:arbitrary` path, the kernel's orthogonality error against the pivoted reference and against the bound `10·n·eps(T)·(1 + ‖Cₖ‖₂²)`, its value against a `BigFloat` solve, and its pullback against a central difference. `test/layers/volume_preserving_attention.jl` is new: 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` at sequence lengths 3 and 7. `test/devices/jlarray.jl` and `test/devices/metal.jl` gain the kernel's device rows, the Metal one a `Float32` testset that also runs the layer's forward pass at N = 2 to 8, and `test/quality/jet.jl` reaches the kernel, its kernel body and its pullback in `Float16`, `Float32` and `Float64`. + `test/kernels/cayley.jl` is new: the exact skewness of the layer's own `:arbitrary` operand, the kernel's orthogonality error against the pivoted reference and against the bound `10·n·eps(T)·(1 + ‖Cₖ‖₂²)`, its value against a `BigFloat` solve, and its pullback against a central difference at every `N` in 2 to 8. `test/layers/volume_preserving_attention.jl` is new: the forward pass of a layer of dimension 4 against an independent host computation and the gradient of the data and of the weight, on the host and on a `JLArray`, at every sequence length in 2 to 8 and with `SL = 0` at sequence lengths 3 and 7. `test/devices/jlarray.jl` and `test/devices/metal.jl` gain the kernel's device rows, the `JLArray` one the same strict rows at every `N` in 2 to 8, the Metal one a `Float32` testset that also runs the layer's forward pass at N = 2 to 8, and `test/quality/jet.jl` reaches the kernel, its kernel body and its pullback in `Float16`, `Float32` and `Float64`. `scripts/utilities/cayley_timings.jl` is new. It times `tensor_cayley` on a batch of 1000 skew slices at scale 1 against the closed form `tensor_cayleyN` of the same `N`, for `N` = 2 to 8 and both element types. Each variant, `N` and element type runs in its own fresh process, with BLAS pinned to one thread and one untimed warm-up call; a time is the median of five calls, forward and backward. The script prints the kernel's two times and the ratio of the kernel to the closed form. For `N` = 6 to 8 there is no closed form in the tree, so it prints the kernel's times alone. - Check 1, the medians the decision rests on, was run at commit `76cee621`, the commit of this branch that still held every closed form, including the batch inverse of `cpu_tensor_cayley`, which the N = 6 to 8 ratios are against. The shipped script cannot reproduce that column, because `cpu_tensor_cayley` is deleted. Machine: Apple M-series, Julia 1.13.2, shared with other work, at a load average between 49 and 88 at the start and the end of the run; BLAS on one thread; batch of 1000 skew slices at scale 1. The two µs columns are the kernel's own forward and backward times, and each ratio is the kernel over the closed form. + Check 1's CPU medians were run at commit `76cee621`, the commit of this branch that still held every closed form, including the batch inverse of `cpu_tensor_cayley`, which the N = 6 to 8 ratios are against. The shipped script cannot reproduce that column, because `cpu_tensor_cayley` is deleted; the Metal column below it is the shipped script's own, as it says there. Machine: Apple M-series, Julia 1.13.2, shared with other work, at a load average between 49 and 88 at the start and the end of the run; BLAS on one thread; batch of 1000 skew slices at scale 1. The two µs columns are the kernel's own forward and backward times, and each ratio is the kernel over the closed form. | N | T | forward kernel µs | backward kernel µs | ratio forward | ratio backward | |:--|:--|--:|--:|--:|--:| @@ -137,17 +137,21 @@ breaking release). then 1.23. Every N from 2 to 5 has a forward ratio above 1 in both runs; every backward ratio on the CPU is below 1 in both runs. - The Metal column, `Float32`, with every timed call wrapped in `Metal.@sync` and the same `measure`, - on an Apple M4 Max (`AGXG16CDevice`), through the test environment, which carries Metal where - `scripts/Project.toml` does not: the ratio of the kernel to the closed form is 0.61 forward and - 5.56 backward at N = 2, 0.41 and 0.71 at N = 3, 1.33 and 1.53 at N = 4, and 1.78 and 0.95 at - N = 5; the kernel's own forward time at N = 6, 7 and 8 is 2544.8, 1771.5 and 2297.3 µs, against - 292.3 at N = 2. The N = 2 backward ratio of 5.56 is the row check 1's rule names for the GPU: the - closed form for N = 2 stays on the CPU columns alone, and on this column too. - - The orthogonality check in `test/kernels/cayley.jl` compares the kernel's `ΛᵀΛ − I` with the pivoted reference `(I + Cₖ) \ (I − Cₖ)` over `N` = 2 to 8, both element types and two scales, on three input sets. The worst ratio of the kernel's error to the reference's, against a limit of 4, is 1.54 for a hand-built skew tensor `R − Rᵀ`, 2.45 for the layer's `:arbitrary` path and 1.30 for its `:skew_sym` path. On the two exactly skew sets the kernel's own error is within the bound `10·n·eps(T)·(1 + ‖Cₖ‖₂²)`, its worst at 0.100 of the bound; on `:skew_sym`, where `C` is skew only to rounding, the pivoted reference itself exceeds that bound by a factor of 3.2, so that set carries the ratio clause alone. - - The closed forms `tensor_cayley2` to `tensor_cayley5` and their dispatch stay. A closed form stays while the kernel is slower than it at a median ratio above 1, and N = 6 to 8 are replaced whatever their ratio. N = 2 stays on its Float64 forward ratio of 1.46. N = 3 stays on its Float32 forward ratio of 1.71 and its Float64 forward ratio of 1.23. N = 4 stays on its Float64 forward ratio of 1.13 in both runs and its Metal ratios of 1.33 forward and 1.53 backward. N = 5 stays on its Float32 forward ratio of 1.43 then 1.42, its Float64 forward ratio of 1.27 then 1.23 and its Metal forward ratio of 1.78. N = 4 and 5 are the user's call rather than the part's, because the kernel is slower forward there in every `Float64` run and on Metal; the user kept their closed forms, `inverse_4x4.jl` and `inverse_5x5.jl` with their 690 generated lines. N = 6 to 8 are replaced: the kernel's ratios there run from 0.16 to 0.64. + The Metal column, `Float32`, is this script's own `measure`, with `MtlArray` as the constructor of + its batch and its seed and with every timed call wrapped in a call that synchronizes the device, + which is what `Metal.@sync` does to an expression; the header documents that call. On an Apple M4 + Max (`AGXG16CDevice`), through the test environment, which carries Metal where + `scripts/Project.toml` does not, the ratio of the kernel to the closed form is 0.73 forward and + 0.84 backward at N = 2, 0.82 and 0.94 at N = 3, 1.51 and 0.90 at N = 4, and 2.28 and 0.93 at N = 5; + a second run on the same machine gives 0.68 and 0.61, 0.76 and 0.99, 1.50 and 0.85, and 2.29 and + 0.99. The kernel's own forward time at N = 6, 7 and 8 is 1582.4, 2103.2 and 2374.5 µs in the first + run and 2806.5, 1845.8 and 1744.0 in the second, against 263.8 and 263.6 at N = 2: the times move + between runs more than the ratios do. The Metal forward ratio is above 1 at N = 4 and N = 5 in both + runs, and below 1 at N = 2 and N = 3. + + The orthogonality check in `test/kernels/cayley.jl` compares the kernel's `‖ΛₖᵀΛₖ − I‖∞` with the pivoted reference `(I + Cₖ) \ (I − Cₖ)` over `N` = 2 to 8, both element types and two scales, on three input sets: the matrix infinity norm of each slice, not its largest entry. The worst ratio of the kernel's error to the reference's, against a limit of 4, is 1.69 for a hand-built skew tensor `R − Rᵀ`, 1.90 for the layer's `:arbitrary` path and 1.44 for its `:skew_sym` path. On the two exactly skew sets the kernel's own error is within the bound `10·n·eps(T)·(1 + ‖Cₖ‖₂²)`, its worst at 0.098 of the bound; on `:skew_sym`, where `C` is skew only to rounding, the pivoted reference itself exceeds that bound by a factor of 6.8, so that set carries the ratio clause alone. The value rows compare each `Λₖ` with a `BigFloat` solve rounded to `T` on the same three sets. + + The closed forms `tensor_cayley2` to `tensor_cayley5` and their dispatch stay. A closed form stays while the kernel is slower than it at a median ratio above 1, and N = 6 to 8 are replaced whatever their ratio. N = 2 stays on its Float64 forward ratio of 1.46. N = 3 stays on its Float32 forward ratio of 1.71 and its Float64 forward ratio of 1.23. N = 4 stays on its Float64 forward ratio of 1.13 in both runs and its Metal forward ratio of 1.51 then 1.50. N = 5 stays on its Float32 forward ratio of 1.43 then 1.42, its Float64 forward ratio of 1.27 then 1.23 and its Metal forward ratio of 2.28 then 2.29. N = 4 and 5 are the user's call rather than the part's, because the kernel is slower forward there in every `Float64` run and on Metal; the user kept their closed forms, `inverse_4x4.jl` and `inverse_5x5.jl` with their 690 generated lines. N = 6 to 8 are replaced: the kernel's ratios there run from 0.16 to 0.64. - `Pkg.test()` no longer runs the doctests. `test/quality/doctests.jl` is now the group `doctests`, which an empty `ARGS` does not run; `Pkg.test(test_args = ["doctests"])` runs it. In CI the Doctests job stays their runner, so the test matrix no longer runs them a second time. - **CI uploads coverage from the `Julia 1 - ubuntu-latest` job instead of `Julia min`, and a test job saves the Julia cache only when it succeeds.** diff --git a/KNOWN_ISSUES.md b/KNOWN_ISSUES.md index 9d53fad06..27f1c4f84 100644 --- a/KNOWN_ISSUES.md +++ b/KNOWN_ISSUES.md @@ -237,6 +237,31 @@ What is known to be wrong in GeometricMachineLearning and is not fixed. Closing it means an error for `size(x, 1) != 2N`. That changes the public behaviour of `assign_q_and_p` and of `PoissonTensor * v`, so it is a change of its own and not a cleanup. +### K19 · In `Float32` the Cayley kernel returns non-finite slices for a skew `C` of norm 1e8, where no method is accurate. + +- **location:** `src/kernels/cayley.jl` +- **kind:** found late +- **found:** 2026-10-17, by the critics of part M4 +- **evidence:** + On 200 skew `Float32` slices of `N = 5` whose largest `‖Cₖ‖₂` is 6.55e8, `tensor_cayley` returns + non-finite values in 23 of them and raises no error. The scale is far outside check 3's `{1, 1e2}`, + and no method is accurate there: `tensor_cayley5`, the closed form that stays, is non-finite in 197 + of the same 200 slices, and the dense solve `(I + Cₖ) \ (I − Cₖ)` is finite on all 200 but has an + orthogonality error `‖ΛₖᵀΛₖ − I‖∞` of 7.3e4 at worst, which is no answer at that scale. What leaves + the `Float32` range is the elimination's own intermediate values rather than a small pivot: on the + two of those slices checked, the same elimination in `Float64`, whose range holds them, reaches a + largest intermediate entry of 4.6e183 and 2.5e207, against a smallest pivot of 1.2e8 and 2.6e8, and + `Float32`'s largest finite value is 3.4e38. The kernel's error on its 177 finite slices is at most + 284. + + ```julia + R = 1.0f8 * randn(Float32, 5, 5, 200) + C = R - permutedims(R, (2, 1, 3)) + nonfinite(L) = count(k -> !all(isfinite, L[:, :, k]), 1:200) + nonfinite(tensor_cayley(C)) # 23 of 200 + nonfinite(tensor_cayley5(C)) # 197 of 200 + ``` + ## Upstream ### K2 · Revise prints EMFILE errors in the test log diff --git a/scripts/utilities/cayley_timings.jl b/scripts/utilities/cayley_timings.jl index 543e729cb..4d121adc4 100644 --- a/scripts/utilities/cayley_timings.jl +++ b/scripts/utilities/cayley_timings.jl @@ -6,27 +6,31 @@ # median of five timed calls. The parent prints the ratio of the kernel to the closed form, forward # and backward. # -# The table that decided which closed forms stay carried a column for N = 6 to 8 as well, against the -# batch inverse of a long slice, which the tree no longer defines. That column came from a commit -# that still held it, and the `CHANGELOG.md` entry of the change names the commit. +# The table that decides which closed forms stay has a column for N = 6 to 8 as well, against the +# batch inverse of a long slice, which the tree does not define: that column belongs to the branch +# commit that holds it, which the `CHANGELOG.md` entry of the change names, so the shipped script +# prints the kernel's own time there instead. # # Run from the repository root, with the package's own environment, which carries `Zygote` and the # standard libraries this script uses; `--project=scripts` resolves them as well, and CUDA with them: # julia --startup-file=no --project=. scripts/utilities/cayley_timings.jl # julia --startup-file=no --project=. scripts/utilities/cayley_timings.jl tensor_cayley 4 Float32 # -# Metal has no reference for N = 6 to 8: the batch inverse that served those rows called `inv` inside -# a kernel, and the tree does not define it, so those rows print the kernel's time alone. N = 2 to 5 -# compare against the closed forms, which run on Metal as well; `Metal` is not a dependency of the -# package, so its column runs `measure` through Kaimon, in a scratch environment that develops this -# worktree and adds Metal. This file runs its table only when it is not included, so a session -# reaches `measure` by including it and calling it: +# Metal has no reference for N = 6 to 8: the batch inverse they would compare against calls `inv` +# inside a kernel, and the tree does not define it, so those rows print the kernel's time alone. +# N = 2 to 5 compare against the closed forms, which run on Metal as well; `Metal` is not a +# dependency of the package, so its column runs `measure` in a process whose environment carries +# Metal, which `scripts/Project.toml` does not. This file runs its table only as the program, so a +# session reaches `measure` by including it and calling it with the constructor that moves the batch +# and the seed onto the device and with the wrapper that makes a timed call a time rather than a +# launch: # # julia> include("scripts/utilities/cayley_timings.jl") -# julia> measure("tensor_cayley", 6, Float32, Metal.@sync) +# julia> sync(x) = (Metal.synchronize(); x) +# julia> measure("tensor_cayley", 6, Float32, sync, MtlArray) # -# `Metal.@sync` is what makes the figure a time rather than a launch: without it the call returns -# before the kernel has run. +# `Metal.synchronize` is what makes the figure a time rather than a launch: without it the call +# returns before the kernel has run. using GeometricMachineLearning using LinearAlgebra: BLAS @@ -40,8 +44,8 @@ const RUNS = 5 median(v) = sort(v)[(length(v) + 1) ÷ 2] -# One untimed warm-up call, then `RUNS` timed calls, in µs. `sync` wraps a timed call and is the -# identity except on Metal. +# One untimed warm-up call, then `RUNS` timed calls, in µs. `sync` wraps one timed call and is the +# identity except on a device, where the figure is a time rather than a launch. function times(f, sync = identity) f() median([(@elapsed sync(f())) * 1e6 for _ in 1:RUNS]) @@ -52,17 +56,18 @@ function skew_batch(::Type{T}, n) where {T} R - permutedims(R, (2, 1, 3)) end -# The forward and the backward time of one variant in this process, in µs. -function measure(variant, n, ::Type{T}, sync = identity) where {T} +# The forward and the backward time of one variant in this process, in µs. `array` moves the batch +# and the seed to the backend the variant runs on, and is `Array` on the CPU. +function measure(variant, n, ::Type{T}, sync = identity, array = Array) where {T} BLAS.set_num_threads(1) Random.seed!(0) - C = skew_batch(T, n) + C = array(skew_batch(T, n)) f = getfield(GeometricMachineLearning, Symbol(variant)) forward = times(() -> f(C), sync) Λ, pullback = Zygote.pullback(f, C) - Λ̄ = rand(T, size(Λ)) + Λ̄ = array(rand(T, size(Λ))) backward = times(() -> pullback(Λ̄), sync) forward, backward @@ -96,8 +101,11 @@ function table() end end -if length(ARGS) == 3 - println(join(measure(ARGS[1], parse(Int, ARGS[2]), getfield(Base, Symbol(ARGS[3]))), " ")) -elseif !isinteractive() - table() +if abspath(PROGRAM_FILE) == abspath(@__FILE__) + if length(ARGS) == 3 + variant, n, T = ARGS[1], parse(Int, ARGS[2]), getfield(Base, Symbol(ARGS[3])) + println(join(measure(variant, n, T), " ")) + else + table() + end end diff --git a/src/kernels/cayley.jl b/src/kernels/cayley.jl index 6d4f79fa3..9a97601ac 100644 --- a/src/kernels/cayley.jl +++ b/src/kernels/cayley.jl @@ -71,13 +71,11 @@ end end if pivot != col - for j in 1:n - a = A[col, j, k] - A[col, j, k] = A[pivot, j, k] - A[pivot, j, k] = a - b = Λ[col, j, k] - Λ[col, j, k] = Λ[pivot, j, k] - Λ[pivot, j, k] = b + for X in (A, Λ), j in 1:n + + a = X[col, j, k] + X[col, j, k] = X[pivot, j, k] + X[pivot, j, k] = a end end diff --git a/src/layers/volume_preserving_attention.jl b/src/layers/volume_preserving_attention.jl index 80e074a92..64ee6e00c 100644 --- a/src/layers/volume_preserving_attention.jl +++ b/src/layers/volume_preserving_attention.jl @@ -24,8 +24,9 @@ The activation is one kernel, [`tensor_cayley`](@ref), which solves ``(\mathbb{I \mathbb{I} - C`` by row-pivoted Gauss–Jordan elimination, one work item per slice: it runs at any sequence length, on the CPU and on a device, and in the element type of the data. -The sequence lengths 2, 3, 4 and 5 use the closed forms `tensor_cayley2` to `tensor_cayley5` -instead, whose inverses have been computed symbolically. +The layer's `seq_length` chooses the activation: a `seq_length` of 2, 3, 4 or 5 uses the closed form +`tensor_cayley2` to `tensor_cayley5`, whose inverse has been computed symbolically, and every other +`seq_length`, `0` included, uses the kernel, whatever the length of the data the layer is applied to. """ struct VolumePreservingAttention{M, N, ScalarProductType, SL} <: AbstractExplicitLayer{M, N} function VolumePreservingAttention(dim::Int, seq_length::Int = 0; skew_sym::Bool = false) @@ -105,12 +106,19 @@ end function (d::VolumePreservingAttention{M, M, :arbitrary})( x::AbstractArray{ T, 3}, ps::NamedTuple) where {T, M} - x_interim = tensor_mat_skew_sym_assign(x, ps.A) / T(√M) - tensor_tensor_mul(x, - orthonormal_activation_cayley(d, x_interim - permutedims(x_interim, (2, 1, 3))) - ) + tensor_tensor_mul(x, orthonormal_activation_cayley(d, arbitrary_operand(x, ps))) end function (d::VolumePreservingAttention)(x::AbstractMatrix, ps::NamedTuple) reshape(d(reshape(x, size(x)..., 1), ps), size(x)...) end + +# The operand `C` of the `:arbitrary` path, whose Cayley transform the layer takes: the scalar +# products of the columns of `x` weighted by `A`, divided by `√M`, made skew by one subtraction of a +# permutation of the same array. Both terms are the same array, so `C + Cᵀ` is exactly zero. The +# `:skew_sym` path forms `xᵀAx` instead, which is skew only to rounding. +function arbitrary_operand(x::AbstractArray{T, 3}, ps::NamedTuple) where {T} + M = size(x, 1) + x_interim = tensor_mat_skew_sym_assign(x, ps.A) / T(√M) + x_interim - permutedims(x_interim, (2, 1, 3)) +end diff --git a/test/devices/jlarray.jl b/test/devices/jlarray.jl index eeb9cf047..82c293798 100644 --- a/test/devices/jlarray.jl +++ b/test/devices/jlarray.jl @@ -62,15 +62,18 @@ end end # The Cayley kernel's own device rows: it launches where scalar indexing is off, its result stays on -# the device, and it matches the per-slice solve on the host at the bound of its orthogonality rows. +# the device and matches the per-slice solve on the host at the bound of its orthogonality rows, and +# its pullback is the derivative of that value. @testset "tensor_cayley and its pullback on a JLArray" begin - for T in (Float32, Float64), n in (2, 3, 8) + for T in (Float32, Float64), n in 2:8 R = rand(T, n, n, 4) C = R - permutedims(R, (2, 1, 3)) + seed = rand(T, n, n, 4) + D = rand(T, n, n, 4) + Λ, pullback = rrule(tensor_cayley, JLArray(C)) - Λ̄ = rand(T, n, n, 4) - _, C̄ = pullback(JLArray(Λ̄)) + _, C̄ = pullback(JLArray(seed)) @test ondevice(Λ) && ondevice(C̄) Id = Matrix{T}(I, n, n) @@ -78,7 +81,20 @@ end bound = 10 * n * eps(T) * (1 + maximum(opnorm(C[:, :, k], 2)^2 for k in axes(C, 3))) @test maximum(abs, Array(Λ) - reference) <= bound - _, C̄_host = rrule(tensor_cayley, C)[2](Λ̄) + # The strict rows of check 4, on the device: the `Float64` central difference of the device + # pullback with `h = cbrt(eps)`, and the `Float32` pullback against the `Float64` one, which + # is compared relatively because its scale is the data's. + if T === Float64 + h = cbrt(eps(T)) + loss(X) = sum(seed .* Array(tensor_cayley(JLArray(X)))) + fd = (loss(C + h * D) - loss(C - h * D)) / (2h) + @test abs(sum(Array(C̄) .* D) - fd) / abs(fd) <= 1e-6 + else + C̄64 = Array(rrule(tensor_cayley, JLArray(Float64.(C)))[2](JLArray(Float64.(seed)))[2]) + @test maximum(abs, Array(C̄) - Float32.(C̄64)) <= bound * maximum(abs, C̄64) + end + + _, C̄_host = rrule(tensor_cayley, C)[2](seed) @test maximum(abs, Array(C̄) - C̄_host) <= bound * maximum(abs, C̄_host) end end diff --git a/test/kernels/cayley.jl b/test/kernels/cayley.jl index 10b374882..e978a4cfd 100644 --- a/test/kernels/cayley.jl +++ b/test/kernels/cayley.jl @@ -8,7 +8,7 @@ # forms are what `tensor_cayley` replaces. using GeometricMachineLearning -using GeometricMachineLearning: tensor_cayley, tensor_mat_skew_sym_assign, +using GeometricMachineLearning: tensor_cayley, arbitrary_operand, tensor_transpose_tensor_mul, _mul using ChainRulesCore: AbstractZero, ZeroTangent, rrule using LinearAlgebra: BLAS, I, opnorm @@ -17,25 +17,15 @@ import Random Random.seed!(1234) -# The references below are solves of 8×8 systems, thousands of them, and threaded LAPACK spends more -# time waking its threads than solving: with BLAS on this machine's twelve threads the orthogonality -# rows over their 84 batches take about two minutes, and 2.6 s on one thread, where the value rows -# take 10.9 s. The file pins one thread for its own work and restores the count it found at the end, -# so no later testset inherits it. +# The references below are solves of up to 8×8 systems, thousands of them, and threaded LAPACK spends +# more time waking its threads than solving: the file pins one thread for its own work and restores +# the count it found at the end, so no later testset inherits it. const BLAS_THREADS = BLAS.get_num_threads() BLAS.set_num_threads(1) const SLICES = 1000 const SCALES = (1, 1e2) -function identity_tensor(::Type{T}, n::Int, k::Int) where {T} - E = zeros(T, n, n, k) - for i in 1:n - E[i, i, :] .= one(T) - end - E -end - # The largest `‖Cₖ‖₂` over the slices of `C`, as a `Float64`. function largest_norm(C::AbstractArray) maximum(Float64(opnorm(@view(C[:, :, k]), 2)) for k in axes(C, 3)) @@ -50,18 +40,14 @@ function skew_tensor(::Type{T}, n::Int, s, k::Int = SLICES) where {T} R - permutedims(R, (2, 1, 3)) end -# (ii) The `:arbitrary` path of the layer: `tensor_mat_skew_sym_assign(x, A) / √M`, then -# `x_interim - permutedims(x_interim, (2, 1, 3))`. That path is linear in `A`, so the whole tensor is -# scaled by one factor on `A` to put its largest `‖Cₖ‖₂` at `s`. +# (ii) The `:arbitrary` path of the layer, which is its own `arbitrary_operand`: +# `tensor_mat_skew_sym_assign(x, A) / √M`, then `x_interim - permutedims(x_interim, (2, 1, 3))`. The +# path is linear in `A`, so scaling the tensor it returns by one factor puts its largest `‖Cₖ‖₂` at +# `s`, which is all `scale_to` needs. function arbitrary_tensor(::Type{T}, n::Int, s) where {T} x = randn(T, n, n, SLICES) A = randn(T, n, n) - scale_to(arbitrary_path(x, A), s) -end - -function arbitrary_path(x::AbstractArray{T, 3}, A::AbstractMatrix) where {T} - x_interim = tensor_mat_skew_sym_assign(x, A) / T(√size(x, 1)) - x_interim - permutedims(x_interim, (2, 1, 3)) + scale_to(arbitrary_operand(x, (A = A,)), s) end # (iii) The `:skew_sym` path of the layer: `Zᵀ A Z` for a `SkewSymMatrix` `A`. That is quadratic in @@ -92,14 +78,21 @@ function big_reference_solve(::Type{T}, Ck::AbstractMatrix) where {T} T.((Id + C) \ (Id - C)) end -# The same maximum over the pivoted references, which is what the ratio clause compares against. +# The orthogonality error of one slice, `‖ΛᵀΛ − I‖∞`: the matrix infinity norm, which is the norm the +# kernel's bound and the ratio clause are written in. +function orthogonality_error(Λ::AbstractMatrix{T}) where {T} + Id = Matrix{T}(I, size(Λ)) + Float64(opnorm(Λ' * Λ - Id, Inf)) +end + +# The batch maximum of that error over the pivoted references, which is what the ratio clause +# compares against. function reference_orthogonality_error(C::AbstractArray{T, 3}) where {T} - n = size(C, 1) - E = Matrix{T}(I, n, n) + # Every slice is measured with `opnorm(ΛᵀΛ − I, Inf)`, the metric the kernel's own error is + # measured in, so the ratio is a ratio of like quantities. worst = 0.0 for k in axes(C, 3) - X = reference_solve(T, @view(C[:, :, k])) - worst = max(worst, Float64(maximum(abs, X' * X - E))) + worst = max(worst, orthogonality_error(reference_solve(T, @view(C[:, :, k])))) end worst end @@ -113,57 +106,58 @@ end # Every row of the orthogonality check: the ratio against the pivoted reference on all three input # sets, and the absolute bound on the two exactly skew ones. function orthogonality_rows(::Type{T}, n::Int, s) where {T} - kernel_error = Ref(0.0) - reference_error = Ref(0.0) - worst_absolute = Ref(0.0) for (set, C) in (("hand-built", skew_tensor(T, n, s)), ("arbitrary", arbitrary_tensor(T, n, s)), ("skew_sym", skew_sym_tensor(T, n, s))) Λ = tensor_cayley(C) - E = identity_tensor(T, n, size(C, 3)) - ΛᵀΛ = tensor_transpose_tensor_mul(Λ, Λ) - kernel_error[] = Float64(maximum(abs, ΛᵀΛ - E)) - reference_error[] = reference_orthogonality_error(C) + kernel_error = maximum(orthogonality_error(@view Λ[:, :, k]) for k in axes(C, 3)) + reference_error = reference_orthogonality_error(C) # The ratio clause holds on every input set; the reference's own error is not zero, so the # factor is a real bound rather than an equality. - @test kernel_error[] <= 4 * reference_error[] + @test kernel_error <= 4 * reference_error + # The absolute clause is per slice; `:skew_sym` carries the ratio clause alone, because its + # own reference breaks the absolute bound there. if set != "skew_sym" + worst = 0.0 for k in axes(C, 3) - bound = bound_per_slice(T, @view(C[:, :, k]), n) - error_k = Float64(maximum(abs, (@view ΛᵀΛ[:, :, k]) .- (@view E[:, :, k]))) - worst_absolute[] = max(worst_absolute[], error_k / bound) + error_k = orthogonality_error(@view Λ[:, :, k]) + worst = max(worst, error_k / bound_per_slice(T, @view(C[:, :, k]), n)) end - @test worst_absolute[] <= 1 + @test worst <= 1 end end end # The value row: per slice, the distance from the `BigFloat` solve, rounded to `T`, under the same -# bound. `Λᵀ` is orthogonal wherever `Λ` is, so the orthogonality rows alone do not pin the value. +# bound, on every input set. `Λᵀ` is orthogonal wherever `Λ` is, so the orthogonality rows alone do +# not pin the value. function value_rows(::Type{T}, n::Int, s) where {T} - for (set, C) in (("hand-built", skew_tensor(T, n, s)), - ("arbitrary", arbitrary_tensor(T, n, s))) + for (_, C) in (("hand-built", skew_tensor(T, n, s)), + ("arbitrary", arbitrary_tensor(T, n, s)), + ("skew_sym", skew_sym_tensor(T, n, s))) Λ = tensor_cayley(C) - worst = Ref(0.0) + worst = 0.0 for k in axes(C, 3) bound = bound_per_slice(T, @view(C[:, :, k]), n) error_k = Float64(maximum(abs, (@view Λ[:, :, k]) .- big_reference_solve(T, @view(C[:, :, k])))) - worst[] = max(worst[], error_k / bound) + worst = max(worst, error_k / bound) end - @test worst[] <= 1 + @test worst <= 1 end end @testset "tensor_cayley: the input of the layer is always skew" begin - for T in (Float32, Float64), n in 2:4 + for T in (Float32, Float64), n in (2, 3, 5) x = randn(T, n, n, 7) A = randn(T, n, n) - # The `:arbitrary` path is exactly skew: one subtraction of a permutation of the same array. - C = arbitrary_path(x, A) + # The layer's own `:arbitrary` operand is exactly skew: one subtraction of a permutation of + # the same array. `M = 4` is the one size here whose `√M` is exact, so the sizes that run + # exercise the rounding of the divisor as well. + C = arbitrary_operand(x, (A = A,)) @test iszero(C + permutedims(C, (2, 1, 3))) # The `:skew_sym` path is skew only to rounding, so nothing is asserted exactly on it. Its # deviation is far below the orthogonality rows' bound, which is what those rows measure. @@ -180,8 +174,7 @@ end value_rows(T, n, s) end -@testset "tensor_cayley: the pullback, N = $n" for T in (Float32, Float64), - n in (2, 3, 5, 8) +@testset "tensor_cayley: the pullback, N = $n" for T in (Float32, Float64), n in 2:8 C = skew_tensor(T, n, 1, 7) Λ̄ = randn(T, n, n, 7) diff --git a/test/layers/volume_preserving_attention.jl b/test/layers/volume_preserving_attention.jl index d082d7884..cb3d06636 100644 --- a/test/layers/volume_preserving_attention.jl +++ b/test/layers/volume_preserving_attention.jl @@ -80,22 +80,23 @@ function tolerance(d, x, ps) 10 * size(x, 2) * eps(T) * (1 + worst) * Float64(maximum(abs, Array(x))) end -# The layer, its parameters, and the data of one case, all on the host; the device case converts -# them, as `test/devices/jlarray.jl` does, rather than initialising them on the device. -function case(N, T; skew_sym) - d = VolumePreservingAttention(N, N; skew_sym) +# A layer of dimension 4 and sequence length `SL` (`0` for the default, which takes data of any +# length), its parameters and its data, all on the host; a device case converts them, as +# `test/devices/jlarray.jl` does, rather than initialising them on the device. +function case(SL, seq_length, T; skew_sym) + d = VolumePreservingAttention(4, SL; skew_sym) model = Chain(d) ps = NeuralNetwork(model, CPU(), T).params.L1 - x = randn(T, N, N, 3) + x = randn(T, 4, seq_length, 3) d, ps, x end @testset "VolumePreservingAttention runs at every sequence length, T = $T" for T in (Float32, Float64) for N in 2:8, skew_sym in (false, true) - d, ps, x = case(N, T; skew_sym) - device = VolumePreservingAttention(N, N; skew_sym) - ps_device = (A = skew_sym ? SkewSymMatrix(JLArray(ps.A.S), N) : JLArray(ps.A),) + d, ps, x = case(N, N, T; skew_sym) + device = VolumePreservingAttention(4, N; skew_sym) + ps_device = (A = skew_sym ? SkewSymMatrix(JLArray(ps.A.S), 4) : JLArray(ps.A),) y_device = device(JLArray(x), ps_device) @test ondevice(y_device) @@ -108,18 +109,17 @@ end gx_device = gradient(z -> sum(device(z, ps_device)), JLArray(x))[1] gx_host = gradient(z -> sum(d(z, ps)), x)[1] @test ondevice(gx_device) - @test maximum(abs, Array(gx_device) - gx_host) <= - tolerance(d, x, ps) * 10 + @test maximum(abs, Array(gx_device) - gx_host) <= tolerance(d, x, ps) # The gradient with respect to the weight, on the device against the host. gA_device = gradient(A -> sum(device(JLArray(x), (A = A,))), ps_device.A)[1] gA_host = gradient(A -> sum(d(x, (A = A,))), ps.A)[1] if skew_sym @test ondevice(gA_device.S) - @test maximum(abs, Array(gA_device.S) - gA_host.S) <= tolerance(d, x, ps) * 10 + @test maximum(abs, Array(gA_device.S) - gA_host.S) <= tolerance(d, x, ps) else @test ondevice(gA_device) - @test maximum(abs, Array(gA_device) - gA_host) <= tolerance(d, x, ps) * 10 + @test maximum(abs, Array(gA_device) - gA_host) <= tolerance(d, x, ps) end end end @@ -128,10 +128,7 @@ end Float64) for seq_length in (3, 7), skew_sym in (false, true) - d = VolumePreservingAttention(4; skew_sym) - model = Chain(d) - ps = NeuralNetwork(model, CPU(), T).params.L1 - x = randn(T, 4, seq_length, 3) + d, ps, x = case(0, seq_length, T; skew_sym) y = d(x, ps) @test size(y) == size(x) @@ -142,12 +139,29 @@ end @test ondevice(y_device) @test maximum(abs, Array(y_device) - reference_forward(d, x, ps)) <= tolerance(d, x, ps) + + # The default layer has no symbolic inverse, so its pullback runs the kernel's; both + # gradients are compared with the host's, as at a fixed sequence length. + gx_device = gradient(z -> sum(d(z, ps_device)), JLArray(x))[1] + gx_host = gradient(z -> sum(d(z, ps)), x)[1] + @test ondevice(gx_device) + @test maximum(abs, Array(gx_device) - gx_host) <= tolerance(d, x, ps) + + gA_device = gradient(A -> sum(d(JLArray(x), (A = A,))), ps_device.A)[1] + gA_host = gradient(A -> sum(d(x, (A = A,))), ps.A)[1] + if skew_sym + @test ondevice(gA_device.S) + @test maximum(abs, Array(gA_device.S) - gA_host.S) <= tolerance(d, x, ps) + else + @test ondevice(gA_device) + @test maximum(abs, Array(gA_device) - gA_host) <= tolerance(d, x, ps) + end end end @testset "the gradient of VolumePreservingAttention at N = 6 against finite differences" begin T = Float64 - d, ps, x = case(6, T; skew_sym = false) + d, ps, x = case(6, 6, T; skew_sym = false) Λ̄ = randn(T, size(x)) fd = j′vp(central_fdm(5, 1), z -> d(z, ps), Λ̄, x)[1] analytic = gradient(z -> sum(Λ̄ .* d(z, ps)), x)[1] From cffd2549e47ec7f1d19acc100028c66a51a932c6 Mon Sep 17 00:00:00 2001 From: Michael Kraus Date: Sat, 17 Oct 2026 02:02:44 +0200 Subject: [PATCH 07/10] GML part M4: one definition of the Cayley bound, for the four rows that compare against it MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- test/devices/jlarray.jl | 7 +++++-- test/devices/metal.jl | 13 +++++++------ test/helpers/cayley_bound.jl | 8 ++++++++ test/kernels/cayley.jl | 14 ++++++-------- test/layers/volume_preserving_attention.jl | 10 +++++++--- 5 files changed, 33 insertions(+), 19 deletions(-) create mode 100644 test/helpers/cayley_bound.jl diff --git a/test/devices/jlarray.jl b/test/devices/jlarray.jl index 82c293798..34c8e8798 100644 --- a/test/devices/jlarray.jl +++ b/test/devices/jlarray.jl @@ -7,12 +7,15 @@ using GeometricMachineLearning: _mul, tensor_tensor_mul, tensor_cayley, assign_q using ChainRulesCore: rrule using GPUArraysCore: allowscalar using JLArrays: JLArray, JLBackend -using LinearAlgebra: I, opnorm +using LinearAlgebra: I import KernelAbstractions using Zygote: gradient using Test import Random +# The bound `10 n eps(T) (1 + ‖Cₖ‖₂²)` of the kernel's orthogonality and value rows. +include("../helpers/cayley_bound.jl") + Random.seed!(123) allowscalar(false) @@ -78,7 +81,7 @@ end Id = Matrix{T}(I, n, n) reference = cat(((Id + C[:, :, k]) \ (Id - C[:, :, k]) for k in axes(C, 3))...; dims = 3) - bound = 10 * n * eps(T) * (1 + maximum(opnorm(C[:, :, k], 2)^2 for k in axes(C, 3))) + bound = maximum(cayley_bound(T, @view(C[:, :, k]), n) for k in axes(C, 3)) @test maximum(abs, Array(Λ) - reference) <= bound # The strict rows of check 4, on the device: the `Float64` central difference of the device diff --git a/test/devices/metal.jl b/test/devices/metal.jl index 127d0f215..668a1ad2c 100644 --- a/test/devices/metal.jl +++ b/test/devices/metal.jl @@ -15,11 +15,14 @@ using GeometricMachineLearning using GeometricMachineLearning: _mul, tensor_tensor_mul, tensor_cayley import KernelAbstractions -using LinearAlgebra: I, dot, opnorm +using LinearAlgebra: I, dot using Metal using Test import Random +# The bound `10 n eps(T) (1 + ‖Cₖ‖₂²)` of the kernel's orthogonality and value rows. +include("../helpers/cayley_bound.jl") + # The layer's forward pass on the host, slice by slice: the same scalar products the kernel computes # for `C`, a dense solve of the Cayley transform, and the product with the data. It returns the # slices of `C` as well, because the bound of the kernel's orthogonality rows is written in them. @@ -95,8 +98,7 @@ if Metal.functional() Id = Matrix{T}(I, n, n) reference = cat(((Id + C[:, :, k]) \ (Id - C[:, :, k]) for k in axes(C, 3))...; dims = 3) - bound = 10 * n * eps(T) * - (1 + maximum(opnorm(C[:, :, k], 2)^2 for k in axes(C, 3))) + bound = maximum(cayley_bound(T, @view(C[:, :, k]), n) for k in axes(C, 3)) @test maximum(abs, Array(Λ) - reference) <= bound d = VolumePreservingAttention(n, n) @@ -104,9 +106,8 @@ if Metal.functional() x = rand(T, n, n, 8) y = d(MtlArray(x), (A = MtlArray(ps.A),)) expected, C_layer = host_forward(x, ps.A) - layer_bound = 10 * n * eps(T) * - (1 + maximum(opnorm(C_layer[:, :, k], 2)^2 for k in axes(x, 3))) * - maximum(abs, x) + layer_bounds = (cayley_bound(T, @view(C_layer[:, :, k]), n) for k in axes(x, 3)) + layer_bound = maximum(layer_bounds) * maximum(abs, x) @test KernelAbstractions.get_backend(y) == backend @test maximum(abs, Array(y) - expected) <= layer_bound end diff --git a/test/helpers/cayley_bound.jl b/test/helpers/cayley_bound.jl new file mode 100644 index 000000000..1d08263b2 --- /dev/null +++ b/test/helpers/cayley_bound.jl @@ -0,0 +1,8 @@ +# The bound of the Cayley kernel's orthogonality and value rows, `10 n eps(T) (1 + ‖Cₖ‖₂²)`, which +# `tensor_cayley`'s docstring states. One definition, because the kernel's own rows, the device rows +# and the layer's forward pass all compare against it and must stay in step with that docstring. +using LinearAlgebra: opnorm + +function cayley_bound(::Type{T}, Ck::AbstractMatrix, n::Int) where {T} + 10.0 * n * eps(T) * (1.0 + Float64(opnorm(Ck, 2))^2) +end diff --git a/test/kernels/cayley.jl b/test/kernels/cayley.jl index e978a4cfd..e152c090c 100644 --- a/test/kernels/cayley.jl +++ b/test/kernels/cayley.jl @@ -15,6 +15,10 @@ using LinearAlgebra: BLAS, I, opnorm using Test import Random +# The bound `10 n eps(T) (1 + ‖Cₖ‖₂²)` of the orthogonality and value rows, shared with the device +# rows of the kernel and with the layer's forward pass. +include("../helpers/cayley_bound.jl") + Random.seed!(1234) # The references below are solves of up to 8×8 systems, thousands of them, and threaded LAPACK spends @@ -97,12 +101,6 @@ function reference_orthogonality_error(C::AbstractArray{T, 3}) where {T} worst end -# The bounds of the orthogonality rows, per slice: the absolute clause `10 n eps(T) (1 + ‖Cₖ‖₂²)`, -# and the ratio clause's own factor on the reference's error. -function bound_per_slice(::Type{T}, Ck::AbstractMatrix, n::Int) where {T} - 10.0 * n * eps(T) * (1.0 + Float64(opnorm(Ck, 2))^2) -end - # Every row of the orthogonality check: the ratio against the pivoted reference on all three input # sets, and the absolute bound on the two exactly skew ones. function orthogonality_rows(::Type{T}, n::Int, s) where {T} @@ -123,7 +121,7 @@ function orthogonality_rows(::Type{T}, n::Int, s) where {T} worst = 0.0 for k in axes(C, 3) error_k = orthogonality_error(@view Λ[:, :, k]) - worst = max(worst, error_k / bound_per_slice(T, @view(C[:, :, k]), n)) + worst = max(worst, error_k / cayley_bound(T, @view(C[:, :, k]), n)) end @test worst <= 1 end @@ -140,7 +138,7 @@ function value_rows(::Type{T}, n::Int, s) where {T} Λ = tensor_cayley(C) worst = 0.0 for k in axes(C, 3) - bound = bound_per_slice(T, @view(C[:, :, k]), n) + bound = cayley_bound(T, @view(C[:, :, k]), n) error_k = Float64(maximum(abs, (@view Λ[:, :, k]) .- big_reference_solve(T, @view(C[:, :, k])))) worst = max(worst, error_k / bound) diff --git a/test/layers/volume_preserving_attention.jl b/test/layers/volume_preserving_attention.jl index cb3d06636..c754be1bc 100644 --- a/test/layers/volume_preserving_attention.jl +++ b/test/layers/volume_preserving_attention.jl @@ -12,11 +12,15 @@ using FiniteDifferences: central_fdm, j′vp using GPUArraysCore: allowscalar using JLArrays: JLArray, JLBackend import KernelAbstractions -using LinearAlgebra: I, dot, opnorm +using LinearAlgebra: I, dot using Zygote: gradient using Test import Random +# The bound `10 n eps(T) (1 + ‖Cₖ‖₂²)` of the kernel's orthogonality and value rows, which the +# forward pass below is compared against, scaled by the data. +include("../helpers/cayley_bound.jl") + Random.seed!(1234) allowscalar(false) @@ -76,8 +80,8 @@ end function tolerance(d, x, ps) T = eltype(x) C = C_tensor(d, x, ps) - worst = maximum(Float64(opnorm(@view(C[:, :, k]), 2))^2 for k in axes(x, 3)) - 10 * size(x, 2) * eps(T) * (1 + worst) * Float64(maximum(abs, Array(x))) + bound = maximum(cayley_bound(T, @view(C[:, :, k]), size(x, 2)) for k in axes(x, 3)) + bound * Float64(maximum(abs, Array(x))) end # A layer of dimension 4 and sequence length `SL` (`0` for the default, which takes data of any From ef1fbae62a5b0739e783a7dfef25e26dc6f38061 Mon Sep 17 00:00:00 2001 From: Michael Kraus Date: Sat, 17 Oct 2026 03:30:32 +0200 Subject: [PATCH 08/10] GML part M4: the verifier's pre-PR fixes, and the entry's check 3 figures `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. --- CHANGELOG.md | 4 ++-- KNOWN_ISSUES.md | 3 +++ scripts/utilities/cayley_timings.jl | 10 ++-------- src/kernels/cayley.jl | 19 ++++++++++++------- src/layers/volume_preserving_attention.jl | 4 ++-- test/kernels/cayley.jl | 16 ++++++++-------- 6 files changed, 29 insertions(+), 27 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index e3d369030..fef57d4bf 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -98,7 +98,7 @@ breaking release). ### Changed -- **`VolumePreservingAttention` computes the Cayley transform in one kernel, `tensor_cayley`, at every sequence length and on a device array; the batch inverse that served the other lengths is deleted, and the closed forms for N = 2 to 5 stay.** +- **`VolumePreservingAttention` computes the Cayley transform on the generic path, the kernel `tensor_cayley`, at every sequence length without a closed form and on a device array; the batch inverse that served those lengths is deleted, and the closed forms for N = 2 to 5 keep their own dispatch.** `src/kernels/cayley.jl` is new. `tensor_cayley(C::AbstractArray{T, 3})` is internal and not exported. It computes `(I + Cₖ)⁻¹(I − Cₖ)` for every slice `k` by row-pivoted Gauss–Jordan elimination on the pair `(I + Cₖ, I − Cₖ)`, one work item per slice, with the loop bound read from `size(C, 1)` at run time. Its `rrule` computes `C̄ = −(I + Λ)ᵀ Λ̄ (I + Λ)ᵀ / 2` from `Λ` alone: no second solve, no identity tensor, and no `Zygote.pullback` of `inv` inside a kernel. The activation of `VolumePreservingAttention` for every sequence length but 2 to 5 runs this kernel, so those lengths run on a device; the path it replaces, `cpu_inverse`, called `inv` and `Zygote.pullback(inv, ·)` on one slice inside a kernel, which this change removes. @@ -149,7 +149,7 @@ breaking release). between runs more than the ratios do. The Metal forward ratio is above 1 at N = 4 and N = 5 in both runs, and below 1 at N = 2 and N = 3. - The orthogonality check in `test/kernels/cayley.jl` compares the kernel's `‖ΛₖᵀΛₖ − I‖∞` with the pivoted reference `(I + Cₖ) \ (I − Cₖ)` over `N` = 2 to 8, both element types and two scales, on three input sets: the matrix infinity norm of each slice, not its largest entry. The worst ratio of the kernel's error to the reference's, against a limit of 4, is 1.69 for a hand-built skew tensor `R − Rᵀ`, 1.90 for the layer's `:arbitrary` path and 1.44 for its `:skew_sym` path. On the two exactly skew sets the kernel's own error is within the bound `10·n·eps(T)·(1 + ‖Cₖ‖₂²)`, its worst at 0.098 of the bound; on `:skew_sym`, where `C` is skew only to rounding, the pivoted reference itself exceeds that bound by a factor of 6.8, so that set carries the ratio clause alone. The value rows compare each `Λₖ` with a `BigFloat` solve rounded to `T` on the same three sets. + The orthogonality check in `test/kernels/cayley.jl` compares the kernel's `‖ΛₖᵀΛₖ − I‖∞` with the pivoted reference `(I + Cₖ) \ (I − Cₖ)` over `N` = 2 to 8, both element types and two scales, on three input sets: the matrix infinity norm of each slice, not its largest entry. The worst ratio of the kernel's error to the reference's, against a limit of 4, is 1.91 for a hand-built skew tensor `R − Rᵀ`, 1.57 for the layer's `:arbitrary` path and 1.34 for its `:skew_sym` path. On the two exactly skew sets the kernel's own error is within the bound `10·n·eps(T)·(1 + ‖Cₖ‖₂²)`, its worst at 0.11 of the bound; on `:skew_sym`, where `C` is skew only to rounding, the pivoted reference itself exceeds that bound by a factor of 3.58, so that set carries the ratio clause alone. The value rows compare each `Λₖ` with a `BigFloat` solve rounded to `T` on the same three sets. The closed forms `tensor_cayley2` to `tensor_cayley5` and their dispatch stay. A closed form stays while the kernel is slower than it at a median ratio above 1, and N = 6 to 8 are replaced whatever their ratio. N = 2 stays on its Float64 forward ratio of 1.46. N = 3 stays on its Float32 forward ratio of 1.71 and its Float64 forward ratio of 1.23. N = 4 stays on its Float64 forward ratio of 1.13 in both runs and its Metal forward ratio of 1.51 then 1.50. N = 5 stays on its Float32 forward ratio of 1.43 then 1.42, its Float64 forward ratio of 1.27 then 1.23 and its Metal forward ratio of 2.28 then 2.29. N = 4 and 5 are the user's call rather than the part's, because the kernel is slower forward there in every `Float64` run and on Metal; the user kept their closed forms, `inverse_4x4.jl` and `inverse_5x5.jl` with their 690 generated lines. N = 6 to 8 are replaced: the kernel's ratios there run from 0.16 to 0.64. diff --git a/KNOWN_ISSUES.md b/KNOWN_ISSUES.md index 27f1c4f84..0278778e4 100644 --- a/KNOWN_ISSUES.md +++ b/KNOWN_ISSUES.md @@ -255,6 +255,9 @@ What is known to be wrong in GeometricMachineLearning and is not fixed. 284. ```julia + using GeometricMachineLearning: tensor_cayley, tensor_cayley5 + import Random + Random.seed!(3) R = 1.0f8 * randn(Float32, 5, 5, 200) C = R - permutedims(R, (2, 1, 3)) nonfinite(L) = count(k -> !all(isfinite, L[:, :, k]), 1:200) diff --git a/scripts/utilities/cayley_timings.jl b/scripts/utilities/cayley_timings.jl index 4d121adc4..ba84d7192 100644 --- a/scripts/utilities/cayley_timings.jl +++ b/scripts/utilities/cayley_timings.jl @@ -6,18 +6,12 @@ # median of five timed calls. The parent prints the ratio of the kernel to the closed form, forward # and backward. # -# The table that decides which closed forms stay has a column for N = 6 to 8 as well, against the -# batch inverse of a long slice, which the tree does not define: that column belongs to the branch -# commit that holds it, which the `CHANGELOG.md` entry of the change names, so the shipped script -# prints the kernel's own time there instead. -# # Run from the repository root, with the package's own environment, which carries `Zygote` and the # standard libraries this script uses; `--project=scripts` resolves them as well, and CUDA with them: # julia --startup-file=no --project=. scripts/utilities/cayley_timings.jl # julia --startup-file=no --project=. scripts/utilities/cayley_timings.jl tensor_cayley 4 Float32 # -# Metal has no reference for N = 6 to 8: the batch inverse they would compare against calls `inv` -# inside a kernel, and the tree does not define it, so those rows print the kernel's time alone. +# Metal rows for N = 6 to 8 print the kernel's time alone. # N = 2 to 5 compare against the closed forms, which run on Metal as well; `Metal` is not a # dependency of the package, so its column runs `measure` in a process whose environment carries # Metal, which `scripts/Project.toml` does not. This file runs its table only as the program, so a @@ -47,7 +41,7 @@ median(v) = sort(v)[(length(v) + 1) ÷ 2] # One untimed warm-up call, then `RUNS` timed calls, in µs. `sync` wraps one timed call and is the # identity except on a device, where the figure is a time rather than a launch. function times(f, sync = identity) - f() + sync(f()) median([(@elapsed sync(f())) * 1e6 for _ in 1:RUNS]) end diff --git a/src/kernels/cayley.jl b/src/kernels/cayley.jl index 9a97601ac..078796b8f 100644 --- a/src/kernels/cayley.jl +++ b/src/kernels/cayley.jl @@ -13,12 +13,16 @@ computed by row-pivoted Gauss–Jordan elimination on the pair ``(\mathbb{I} + C [`VolumePreservingAttention`](@ref), which forms ``C`` skew-symmetric. `C` is assumed skew-symmetric, and this is not checked: a check would cost a reduction per slice on -a device. The result is the Cayley transform of any `C` whose ``\mathbb{I} + C_k`` is nonsingular; -for a skew `C` its orthogonality error ``\|\Lambda_k^T\Lambda_k - \mathbb{I}\|_\infty`` is the error -of a row-pivoted elimination — at most four times that of a dense `(I + C_k) \ (I - C_k)`, and at -most ``10n\varepsilon(1 + \|C_k\|_2^2)`` for the element type's machine epsilon ``\varepsilon``, -where `n = size(C, 1)`. A skew-symmetric ``C_k`` has the positive definite symmetric part -``\mathbb{I}``, which is what bounds the pivots. +a device. For nonsingular ``\mathbb{I} + C_k``, the algorithm computes the Cayley transform. + +The tests exercise sizes 2 to 8 at scales 1 and 1e2. On exactly skew inputs, they bound each +slice's orthogonality error ``\|\Lambda_k^T\Lambda_k - \mathbb{I}\|_\infty`` by +``10n\varepsilon(1 + \|C_k\|_2^2)``, where `n = size(C, 1)` and ``\varepsilon`` is the element +type's machine epsilon. They also bound the batch maximum error by four times the batch maximum +of a dense `(I + C_k) \ (I - C_k)` reference, including the layer's inputs that are skew only to +rounding. These are tested bounds, not guarantees for arbitrary magnitudes: elimination can +overflow at large scales. For skew ``C_k``, the symmetric part of ``\mathbb{I} + C_k`` is +positive definite. The element type `T` is an `AbstractFloat`. Every buffer comes from `similar(C)`, so the result has the element type `T` and lives on the backend of `C`. The loop bound is read from `size(C, 1)` at @@ -108,7 +112,8 @@ function ChainRulesCore.rrule(::typeof(tensor_cayley), C::AbstractArray{T, 3}) w Λ̄ = _seed(T, Λ̄) Λ̄ isa AbstractZero && return NoTangent(), ZeroTangent() Y = Λ̄ + tensor_transpose_tensor_mul(Λ, Λ̄) - NoTangent(), -(Y + tensor_tensor_transpose_mul(Y, Λ)) / 2 + Z = tensor_tensor_transpose_mul(Y, Λ) + NoTangent(), @. -(Y + Z) / T(2) end Λ, tensor_cayley_pullback diff --git a/src/layers/volume_preserving_attention.jl b/src/layers/volume_preserving_attention.jl index 64ee6e00c..e88da15b7 100644 --- a/src/layers/volume_preserving_attention.jl +++ b/src/layers/volume_preserving_attention.jl @@ -20,8 +20,8 @@ Applying a layer of type `VolumePreservingAttention` does the following: # Implementation -The activation is one kernel, [`tensor_cayley`](@ref), which solves ``(\mathbb{I} + C)\Lambda = -\mathbb{I} - C`` by row-pivoted Gauss–Jordan elimination, one work item per slice: it runs at any +The generic activation uses [`tensor_cayley`](@ref), which solves ``(\mathbb{I} + C)\Lambda = +\mathbb{I} - C`` by row-pivoted Gauss–Jordan elimination, one work item per slice. It runs at any sequence length, on the CPU and on a device, and in the element type of the data. The layer's `seq_length` chooses the activation: a `seq_length` of 2, 3, 4 or 5 uses the closed form diff --git a/test/kernels/cayley.jl b/test/kernels/cayley.jl index e152c090c..97d37080d 100644 --- a/test/kernels/cayley.jl +++ b/test/kernels/cayley.jl @@ -4,8 +4,7 @@ # # The three input sets are the ones the layer produces: a hand-built skew tensor, the `:arbitrary` # path and the `:skew_sym` path, both of them through the layer's own functions. The reference is the -# per-slice solve `(I + Cₖ) \ (I - Cₖ)` in the working precision, written here, because the closed -# forms are what `tensor_cayley` replaces. +# per-slice solve `(I + Cₖ) \ (I - Cₖ)` in the working precision, independent of the kernel. using GeometricMachineLearning using GeometricMachineLearning: tensor_cayley, arbitrary_operand, @@ -36,9 +35,9 @@ function largest_norm(C::AbstractArray) end # The tensor scaled so that its largest `‖Cₖ‖₂` is `s`. -scale_to(C::AbstractArray, s) = C * (s / largest_norm(C)) +scale_to(C::AbstractArray{T}, s) where {T} = C * T(s / largest_norm(C)) -# (i) A hand-built skew tensor, `R - Rᵀ` for `R = s * randn`, which is skew to rounding. +# (i) A hand-built skew tensor, `R - Rᵀ` for `R = s * randn`, which is exactly skew. function skew_tensor(::Type{T}, n::Int, s, k::Int = SLICES) where {T} R = T(s) * randn(T, n, n, k) R - permutedims(R, (2, 1, 3)) @@ -67,10 +66,9 @@ function skew_sym_path(x::AbstractArray{T, 3}, A::AbstractMatrix) where {T} tensor_transpose_tensor_mul(x, _mul(A, x)) end -# The per-slice reference of the orthogonality rows and of the value row: the solve -# `(I + Cₖ) \ (I - Cₖ)` in the working precision, and in `BigFloat` for the value row. The two -# agree to the last bits of `T`, so the value row's reference is the transform itself rather than -# another elimination. +# The per-slice reference of the orthogonality rows is the solve `(I + Cₖ) \ (I - Cₖ)` in the +# working precision. The value row solves in `BigFloat` and rounds to `T`, so its reference +# approximates the transform itself rather than another elimination in the working precision. function reference_solve(::Type{T}, Ck::AbstractMatrix) where {T} Id = Matrix{T}(I, size(Ck)) (Id + Ck) \ (Id - Ck) @@ -107,6 +105,7 @@ function orthogonality_rows(::Type{T}, n::Int, s) where {T} for (set, C) in (("hand-built", skew_tensor(T, n, s)), ("arbitrary", arbitrary_tensor(T, n, s)), ("skew_sym", skew_sym_tensor(T, n, s))) + @test eltype(C) === T Λ = tensor_cayley(C) kernel_error = maximum(orthogonality_error(@view Λ[:, :, k]) for k in axes(C, 3)) reference_error = reference_orthogonality_error(C) @@ -135,6 +134,7 @@ function value_rows(::Type{T}, n::Int, s) where {T} for (_, C) in (("hand-built", skew_tensor(T, n, s)), ("arbitrary", arbitrary_tensor(T, n, s)), ("skew_sym", skew_sym_tensor(T, n, s))) + @test eltype(C) === T Λ = tensor_cayley(C) worst = 0.0 for k in axes(C, 3) From 8a4d174c4702b3e021141663c09bcc47a7a600f8 Mon Sep 17 00:00:00 2001 From: Michael Kraus Date: Sat, 17 Oct 2026 10:56:23 +0200 Subject: [PATCH 09/10] Address review of #347: @inbounds Cayley kernel, two-tensor pullback, 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 --- CHANGELOG.md | 82 ++++++++----------- KNOWN_ISSUES.md | 13 +-- src/kernels/cayley.jl | 94 +++++++++------------- src/layers/volume_preserving_attention.jl | 13 ++- test/devices/jlarray.jl | 15 ++-- test/devices/metal.jl | 47 +++-------- test/helpers/cayley.jl | 48 +++++++++++ test/helpers/cayley_bound.jl | 8 -- test/kernels/cayley.jl | 53 +++++------- test/kernels/kernel_pullbacks.jl | 11 +-- test/layers/volume_preserving_attention.jl | 43 ++-------- test/quality/jet.jl | 4 +- 12 files changed, 180 insertions(+), 251 deletions(-) create mode 100644 test/helpers/cayley.jl delete mode 100644 test/helpers/cayley_bound.jl diff --git a/CHANGELOG.md b/CHANGELOG.md index fef57d4bf..a8c2291e7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -98,60 +98,46 @@ breaking release). ### Changed -- **`VolumePreservingAttention` computes the Cayley transform on the generic path, the kernel `tensor_cayley`, at every sequence length without a closed form and on a device array; the batch inverse that served those lengths is deleted, and the closed forms for N = 2 to 5 keep their own dispatch.** +- **`VolumePreservingAttention` computes the Cayley transform with the kernel `tensor_cayley` at every sequence length other than 2 to 5, 0 included, and on device arrays; the batch inverse that served those lengths is deleted.** `tensor_cayley(C)` is internal and not exported. It computes `(I + Cₖ)⁻¹(I − Cₖ)` for every slice `k` by row-pivoted Gauss–Jordan elimination, one work item per slice. Its `rrule` computes `C̄ = −(I + Λ)ᵀ Λ̄ (I + Λ)ᵀ / 2` from `Λ` alone, with no second solve and no identity tensor. - `src/kernels/cayley.jl` is new. `tensor_cayley(C::AbstractArray{T, 3})` is internal and not exported. It computes `(I + Cₖ)⁻¹(I − Cₖ)` for every slice `k` by row-pivoted Gauss–Jordan elimination on the pair `(I + Cₖ, I − Cₖ)`, one work item per slice, with the loop bound read from `size(C, 1)` at run time. Its `rrule` computes `C̄ = −(I + Λ)ᵀ Λ̄ (I + Λ)ᵀ / 2` from `Λ` alone: no second solve, no identity tensor, and no `Zygote.pullback` of `inv` inside a kernel. The activation of `VolumePreservingAttention` for every sequence length but 2 to 5 runs this kernel, so those lengths run on a device; the path it replaces, `cpu_inverse`, called `inv` and `Zygote.pullback(inv, ·)` on one slice inside a kernel, which this change removes. + `cpu_inverse`, its `rrule` and `cpu_tensor_cayley` are deleted. The identity tensor that the closed forms use moved from `src/kernels/exponentials/tensor_exponential.jl` to `src/kernels/identity_tensor.jl`, so that the file name matches what it holds; this closes known issue C18, which leaves `KNOWN_ISSUES.md`. - `cpu_inverse`, its `rrule` and `cpu_tensor_cayley` are deleted, because nothing calls them after this change. The identity tensor that the closed forms use moved from `src/kernels/exponentials/tensor_exponential.jl` to `src/kernels/identity_tensor.jl`, which holds `init_output`, `assign_ones!` and the `rrule` of `init_output`, so the file name matches what it holds, and known issue C18 leaves `KNOWN_ISSUES.md`. + New tests: `test/kernels/cayley.jl` checks the kernel's orthogonality and value, and its pullback against a central difference, for N = 2 to 8. `test/layers/volume_preserving_attention.jl` checks the layer's forward pass and gradients for N = 2 to 8, on the host and on a `JLArray`. `test/devices/jlarray.jl`, `test/devices/metal.jl` and `test/quality/jet.jl` reach the kernel, the last in `Float16`, `Float32` and `Float64`. - `test/kernels/cayley.jl` is new: the exact skewness of the layer's own `:arbitrary` operand, the kernel's orthogonality error against the pivoted reference and against the bound `10·n·eps(T)·(1 + ‖Cₖ‖₂²)`, its value against a `BigFloat` solve, and its pullback against a central difference at every `N` in 2 to 8. `test/layers/volume_preserving_attention.jl` is new: the forward pass of a layer of dimension 4 against an independent host computation and the gradient of the data and of the weight, on the host and on a `JLArray`, at every sequence length in 2 to 8 and with `SL = 0` at sequence lengths 3 and 7. `test/devices/jlarray.jl` and `test/devices/metal.jl` gain the kernel's device rows, the `JLArray` one the same strict rows at every `N` in 2 to 8, the Metal one a `Float32` testset that also runs the layer's forward pass at N = 2 to 8, and `test/quality/jet.jl` reaches the kernel, its kernel body and its pullback in `Float16`, `Float32` and `Float64`. + `scripts/utilities/cayley_timings.jl` is new. It times `tensor_cayley` against the closed form `tensor_cayleyN` on a batch of 1000 skew slices at scale 1, for N = 2 to 8 and both element types. Each variant, N and element type runs in a fresh process with BLAS on one thread and one untimed warm-up call; a time is the median of five calls, forward and backward. For N = 6 to 8 there is no closed form, so the script prints the kernel's times alone. - `scripts/utilities/cayley_timings.jl` is new. It times `tensor_cayley` on a batch of 1000 skew slices at scale 1 against the closed form `tensor_cayleyN` of the same `N`, for `N` = 2 to 8 and both element types. Each variant, `N` and element type runs in its own fresh process, with BLAS pinned to one thread and one untimed warm-up call; a time is the median of five calls, forward and backward. The script prints the kernel's two times and the ratio of the kernel to the closed form. For `N` = 6 to 8 there is no closed form in the tree, so it prints the kernel's times alone. + CPU, from two runs of the script on an Apple M-series machine with Julia 1.13.2, at a load average of 3.7 to 4.6. The ratio is the kernel over the closed form. - Check 1's CPU medians were run at commit `76cee621`, the commit of this branch that still held every closed form, including the batch inverse of `cpu_tensor_cayley`, which the N = 6 to 8 ratios are against. The shipped script cannot reproduce that column, because `cpu_tensor_cayley` is deleted; the Metal column below it is the shipped script's own, as it says there. Machine: Apple M-series, Julia 1.13.2, shared with other work, at a load average between 49 and 88 at the start and the end of the run; BLAS on one thread; batch of 1000 skew slices at scale 1. The two µs columns are the kernel's own forward and backward times, and each ratio is the kernel over the closed form. - - | N | T | forward kernel µs | backward kernel µs | ratio forward | ratio backward | + | N | T | forward µs, run 1 / run 2 | backward µs, run 1 / run 2 | ratio forward, run 1 / run 2 | ratio backward, run 1 / run 2 | |:--|:--|--:|--:|--:|--:| - | 2 | Float32 | 45.3 | 43.0 | 0.47 | 0.63 | - | 3 | Float32 | 74.7 | 72.3 | 1.71 | 0.58 | - | 4 | Float32 | 123.8 | 135.1 | 0.59 | 0.66 | - | 5 | Float32 | 184.5 | 224.3 | 1.43 | 0.66 | - | 6 | Float32 | 606.8 | 337.0 | 0.64 | 0.26 | - | 7 | Float32 | 372.0 | 485.5 | 0.33 | 0.30 | - | 8 | Float32 | 521.4 | 656.9 | 0.38 | 0.33 | - | 2 | Float64 | 48.8 | 51.0 | 1.46 | 0.65 | - | 3 | Float64 | 75.4 | 100.7 | 1.23 | 0.69 | - | 4 | Float64 | 124.4 | 168.6 | 1.13 | 0.62 | - | 5 | Float64 | 185.2 | 245.0 | 1.27 | 0.60 | - | 6 | Float64 | 289.8 | 384.5 | 0.16 | 0.23 | - | 7 | Float64 | 416.1 | 533.4 | 0.30 | 0.28 | - | 8 | Float64 | 565.8 | 756.6 | 0.34 | 0.33 | - - Run again on the shipped script at a lower load average (14.8 to 25.0), with the closed forms it - still holds, the decision is the same. The kernel's own times agree within 3 % of the table above, - except one N = 6 Float32 forward reading, 606.8 → 263.1. The closed form's own time is what moves, - by up to 3.3× (96 µs to 28.8 µs at N = 2, Float32, forward), so one median at these sizes is not a - figure of the code alone. The forward ratios, first run then second: N = 2, Float32 0.47 then 1.51 - and Float64 1.46 then 1.33; N = 3, Float32 1.71 then 1.79 and Float64 1.23 then 1.68; N = 4, - Float32 0.59 then 1.70 and Float64 1.13 in both; N = 5, Float32 1.43 then 1.42 and Float64 1.27 - then 1.23. Every N from 2 to 5 has a forward ratio above 1 in both runs; every backward ratio on - the CPU is below 1 in both runs. - - The Metal column, `Float32`, is this script's own `measure`, with `MtlArray` as the constructor of - its batch and its seed and with every timed call wrapped in a call that synchronizes the device, - which is what `Metal.@sync` does to an expression; the header documents that call. On an Apple M4 - Max (`AGXG16CDevice`), through the test environment, which carries Metal where - `scripts/Project.toml` does not, the ratio of the kernel to the closed form is 0.73 forward and - 0.84 backward at N = 2, 0.82 and 0.94 at N = 3, 1.51 and 0.90 at N = 4, and 2.28 and 0.93 at N = 5; - a second run on the same machine gives 0.68 and 0.61, 0.76 and 0.99, 1.50 and 0.85, and 2.29 and - 0.99. The kernel's own forward time at N = 6, 7 and 8 is 1582.4, 2103.2 and 2374.5 µs in the first - run and 2806.5, 1845.8 and 1744.0 in the second, against 263.8 and 263.6 at N = 2: the times move - between runs more than the ratios do. The Metal forward ratio is above 1 at N = 4 and N = 5 in both - runs, and below 1 at N = 2 and N = 3. - - The orthogonality check in `test/kernels/cayley.jl` compares the kernel's `‖ΛₖᵀΛₖ − I‖∞` with the pivoted reference `(I + Cₖ) \ (I − Cₖ)` over `N` = 2 to 8, both element types and two scales, on three input sets: the matrix infinity norm of each slice, not its largest entry. The worst ratio of the kernel's error to the reference's, against a limit of 4, is 1.91 for a hand-built skew tensor `R − Rᵀ`, 1.57 for the layer's `:arbitrary` path and 1.34 for its `:skew_sym` path. On the two exactly skew sets the kernel's own error is within the bound `10·n·eps(T)·(1 + ‖Cₖ‖₂²)`, its worst at 0.11 of the bound; on `:skew_sym`, where `C` is skew only to rounding, the pivoted reference itself exceeds that bound by a factor of 3.58, so that set carries the ratio clause alone. The value rows compare each `Λₖ` with a `BigFloat` solve rounded to `T` on the same three sets. - - The closed forms `tensor_cayley2` to `tensor_cayley5` and their dispatch stay. A closed form stays while the kernel is slower than it at a median ratio above 1, and N = 6 to 8 are replaced whatever their ratio. N = 2 stays on its Float64 forward ratio of 1.46. N = 3 stays on its Float32 forward ratio of 1.71 and its Float64 forward ratio of 1.23. N = 4 stays on its Float64 forward ratio of 1.13 in both runs and its Metal forward ratio of 1.51 then 1.50. N = 5 stays on its Float32 forward ratio of 1.43 then 1.42, its Float64 forward ratio of 1.27 then 1.23 and its Metal forward ratio of 2.28 then 2.29. N = 4 and 5 are the user's call rather than the part's, because the kernel is slower forward there in every `Float64` run and on Metal; the user kept their closed forms, `inverse_4x4.jl` and `inverse_5x5.jl` with their 690 generated lines. N = 6 to 8 are replaced: the kernel's ratios there run from 0.16 to 0.64. + | 2 | Float32 | 23.9 / 24.2 | 31.0 / 31.0 | 0.85 / 0.86 | 0.51 / 0.49 | + | 3 | Float32 | 54.1 / 54.4 | 62.5 / 64.7 | 1.06 / 1.17 | 0.53 / 0.54 | + | 4 | Float32 | 93.2 / 93.6 | 109.2 / 116.0 | 1.24 / 1.31 | 0.54 / 0.57 | + | 5 | Float32 | 139.7 / 140.5 | 173.6 / 169.7 | 1.07 / 1.04 | 0.50 / 0.50 | + | 6 | Float32 | 209.1 / 214.0 | 293.4 / 280.4 | no closed form | no closed form | + | 7 | Float32 | 315.9 / 307.5 | 429.1 / 410.4 | no closed form | no closed form | + | 8 | Float32 | 407.2 / 406.1 | 545.2 / 553.1 | no closed form | no closed form | + | 2 | Float64 | 27.9 / 28.5 | 31.0 / 34.6 | 0.83 / 0.87 | 0.43 / 0.47 | + | 3 | Float64 | 57.4 / 59.2 | 59.9 / 63.3 | 1.25 / 1.42 | 0.40 / 0.41 | + | 4 | Float64 | 94.6 / 93.8 | 111.8 / 119.2 | 1.06 / 0.93 | 0.42 / 0.46 | + | 5 | Float64 | 145.3 / 145.5 | 184.5 / 178.0 | 1.09 / 0.86 | 0.43 / 0.41 | + | 6 | Float64 | 234.8 / 240.6 | 288.8 / 289.8 | no closed form | no closed form | + | 7 | Float64 | 326.3 / 323.6 | 468.4 / 479.0 | no closed form | no closed form | + | 8 | Float64 | 459.7 / 449.1 | 551.0 / 591.6 | no closed form | no closed form | + + At N = 2 the kernel's forward time is below the closed form's in both runs and both element types, + at ratios 0.83 to 0.87. At N = 3 to 5 the forward ratios lie between 0.86 and 1.42, above 1 in some + runs and below 1 in others. The backward ratio is below 1 at N = 2 to 5 in every run. + + Metal, `Float32`, on the `AGXG16CDevice` Metal device, with the script's own `measure` on `MtlArray` + and every timed call synchronized on the device, as the script's header documents. These are + warm-session figures from four runs. The forward ratio at N = 2 is 0.28, 1.28, 0.86 and 0.51, so one + run's ratio moves by up to 4.6× at N = 2. At N = 3 to 5 the forward ratios lie between 0.48 and 1.20. + The backward ratios lie between 0.44 and 0.99 at every N. + + The orthogonality check in `test/kernels/cayley.jl` compares the kernel's `‖ΛₖᵀΛₖ − I‖∞` with that of the pivoted reference `(I + Cₖ) \ (I − Cₖ)`, for N = 2 to 8, both element types and two scales, on three input sets. The ratio of the kernel's error to the reference's is held to a limit of 4 on all three sets. On the two exactly skew sets the kernel's error must also lie within `10·n·eps(T)·(1 + ‖Cₖ‖₂²)`; on the `:skew_sym` set, where `C` is skew only to rounding, the ratio is the only check. The value rows compare each `Λₖ` with a `BigFloat` solve rounded to `T` on the same three sets. + + The closed forms `tensor_cayley2` to `tensor_cayley5` and their dispatch are kept for sequence lengths 2 to 5. Every other sequence length, 0 included, uses `tensor_cayley` on the CPU and on devices. The figures for N = 2 to 5 are in the tables above. - `Pkg.test()` no longer runs the doctests. `test/quality/doctests.jl` is now the group `doctests`, which an empty `ARGS` does not run; `Pkg.test(test_args = ["doctests"])` runs it. In CI the Doctests job stays their runner, so the test matrix no longer runs them a second time. - **CI uploads coverage from the `Julia 1 - ubuntu-latest` job instead of `Julia min`, and a test job saves the Julia cache only when it succeeds.** diff --git a/KNOWN_ISSUES.md b/KNOWN_ISSUES.md index 0278778e4..85bf93886 100644 --- a/KNOWN_ISSUES.md +++ b/KNOWN_ISSUES.md @@ -247,12 +247,13 @@ What is known to be wrong in GeometricMachineLearning and is not fixed. non-finite values in 23 of them and raises no error. The scale is far outside check 3's `{1, 1e2}`, and no method is accurate there: `tensor_cayley5`, the closed form that stays, is non-finite in 197 of the same 200 slices, and the dense solve `(I + Cₖ) \ (I − Cₖ)` is finite on all 200 but has an - orthogonality error `‖ΛₖᵀΛₖ − I‖∞` of 7.3e4 at worst, which is no answer at that scale. What leaves - the `Float32` range is the elimination's own intermediate values rather than a small pivot: on the - two of those slices checked, the same elimination in `Float64`, whose range holds them, reaches a - largest intermediate entry of 4.6e183 and 2.5e207, against a smallest pivot of 1.2e8 and 2.6e8, and - `Float32`'s largest finite value is 3.4e38. The kernel's error on its 177 finite slices is at most - 284. + orthogonality error `‖ΛₖᵀΛₖ − I‖∞` of 7.3e4 at worst, which is no answer at that scale. The cause + is cancellation in the last pivot, not overflow: in all 23 non-finite slices the last `Float32` + pivot is exactly `0`. The same elimination in `Float64` gives a last pivot of magnitude 1.7 to 15 + on those slices, and no intermediate entry above 6.8e8. In `Float32` the rounding of the earlier + steps, about `eps(Float32) · ‖Cₖ‖₂ ≈ 78`, is larger than that pivot; for odd `N` the skew `Cₖ` is + singular, so the pivot of `I + Cₖ` is small at that scale. The kernel's error on its 177 finite + slices is at most 284. ```julia using GeometricMachineLearning: tensor_cayley, tensor_cayley5 diff --git a/src/kernels/cayley.jl b/src/kernels/cayley.jl index 078796b8f..0020914a9 100644 --- a/src/kernels/cayley.jl +++ b/src/kernels/cayley.jl @@ -9,29 +9,13 @@ is ``` computed by row-pivoted Gauss–Jordan elimination on the pair ``(\mathbb{I} + C_k, -\mathbb{I} - C_k)``, which leaves ``\Lambda_k`` in the second member. This is the transform of -[`VolumePreservingAttention`](@ref), which forms ``C`` skew-symmetric. - -`C` is assumed skew-symmetric, and this is not checked: a check would cost a reduction per slice on -a device. For nonsingular ``\mathbb{I} + C_k``, the algorithm computes the Cayley transform. - -The tests exercise sizes 2 to 8 at scales 1 and 1e2. On exactly skew inputs, they bound each -slice's orthogonality error ``\|\Lambda_k^T\Lambda_k - \mathbb{I}\|_\infty`` by -``10n\varepsilon(1 + \|C_k\|_2^2)``, where `n = size(C, 1)` and ``\varepsilon`` is the element -type's machine epsilon. They also bound the batch maximum error by four times the batch maximum -of a dense `(I + C_k) \ (I - C_k)` reference, including the layer's inputs that are skew only to -rounding. These are tested bounds, not guarantees for arbitrary magnitudes: elimination can -overflow at large scales. For skew ``C_k``, the symmetric part of ``\mathbb{I} + C_k`` is -positive definite. - -The element type `T` is an `AbstractFloat`. Every buffer comes from `similar(C)`, so the result has -the element type `T` and lives on the backend of `C`. The loop bound is read from `size(C, 1)` at -run time, so a slice of any size is solved by the same kernel. - -A slice whose ``\mathbb{I} + C_k`` is singular gives non-finite entries in the slice of that `k` and -no error, as the closed forms do, which divide by the determinant of the slice; a `NaN` or an `Inf` -in `C` spreads within its own slice. A slice that is not square raises a `DimensionMismatch`, and -an empty batch is not reached by a caller of this function. +\mathbb{I} - C_k)``, one work item per slice, for slices of any size. This is the transform of +[`VolumePreservingAttention`](@ref), which forms ``C`` skew-symmetric. `C` is assumed skew and this +is not checked; for skew ``C_k``, ``\mathbb{I} + C_k`` is nonsingular. + +The result is `similar(C)`: it has the element type of `C` and lives on its backend. A singular +``\mathbb{I} + C_k``, or a `NaN` or an `Inf` in ``C_k``, gives non-finite entries in slice ``k`` +and no error. A slice that is not square raises a `DimensionMismatch`. """ function tensor_cayley(C::AbstractArray{T, 3}) where {T} n = size(C, 1) @@ -52,49 +36,43 @@ end # `I + C`, `Λ` with `I - C`, and the row operations that turn `A` into `I` are also applied to `Λ`. # The pivot of a column is the row of the largest magnitude in it, and the swap moves both members: # the two are one augmented system, and a swap on one of them alone solves a permuted system. +# Every index is in bounds: `tensor_cayley` checks that the slices are square, `A` and `Λ` are +# `similar(C)`, and KernelAbstractions masks `k` to the `ndrange`. @kernel function tensor_cayley_kernel!(A, Λ, C) k = @index(Global) n = size(C, 1) - for i in 1:n, j in 1:n - - δ = i == j ? one(eltype(C)) : zero(eltype(C)) - A[i, j, k] = δ + C[i, j, k] - Λ[i, j, k] = δ - C[i, j, k] - end + @inbounds begin + for j in 1:n, i in 1:n - for col in 1:n - pivot = col - largest = abs(A[col, col, k]) - for r in (col + 1):n - candidate = abs(A[r, col, k]) - if candidate > largest - largest = candidate - pivot = r - end + δ = 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 - if pivot != col - for X in (A, Λ), j in 1:n + 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 - a = X[col, j, k] - X[col, j, k] = X[pivot, j, k] - X[pivot, j, k] = a + 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 - end - diagonal = A[col, col, k] - for j in 1:n - A[col, j, k] = A[col, j, k] / diagonal - Λ[col, j, k] = Λ[col, j, k] / diagonal - 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 - factor = A[r, col, k] - if r != col && factor != zero(eltype(C)) + for r in 1:n + r == col && continue + f = A[r, col, k] for j in 1:n - A[r, j, k] = A[r, j, k] - factor * A[col, j, k] - Λ[r, j, k] = Λ[r, j, k] - factor * Λ[col, j, k] + A[r, j, k] -= f * A[col, j, k] + Λ[r, j, k] -= f * Λ[col, j, k] end end end @@ -111,9 +89,11 @@ function ChainRulesCore.rrule(::typeof(tensor_cayley), C::AbstractArray{T, 3}) w function tensor_cayley_pullback(Λ̄) Λ̄ = _seed(T, Λ̄) Λ̄ isa AbstractZero && return NoTangent(), ZeroTangent() - Y = Λ̄ + tensor_transpose_tensor_mul(Λ, Λ̄) - Z = tensor_tensor_transpose_mul(Y, Λ) - NoTangent(), @. -(Y + Z) / T(2) + Y = tensor_transpose_tensor_mul(Λ, Λ̄) + Y .+= Λ̄ + C̄ = tensor_tensor_transpose_mul(Y, Λ) + @. C̄ = -(Y + C̄) / T(2) + NoTangent(), C̄ end Λ, tensor_cayley_pullback diff --git a/src/layers/volume_preserving_attention.jl b/src/layers/volume_preserving_attention.jl index e88da15b7..18a6a4e0c 100644 --- a/src/layers/volume_preserving_attention.jl +++ b/src/layers/volume_preserving_attention.jl @@ -4,8 +4,7 @@ Make an instance of `VolumePreservingAttention` for a specific dimension and sequence length. The sequence length is `0` by default, and a layer built with `0` applies to data of any sequence -length: the Cayley transform of the data is one kernel, [`tensor_cayley`](@ref), rather than a -symbolic inverse, and every sequence length runs on the CPU and on a device. +length. # Arguments @@ -20,13 +19,11 @@ Applying a layer of type `VolumePreservingAttention` does the following: # Implementation -The generic activation uses [`tensor_cayley`](@ref), which solves ``(\mathbb{I} + C)\Lambda = -\mathbb{I} - C`` by row-pivoted Gauss–Jordan elimination, one work item per slice. It runs at any -sequence length, on the CPU and on a device, and in the element type of the data. - The layer's `seq_length` chooses the activation: a `seq_length` of 2, 3, 4 or 5 uses the closed form -`tensor_cayley2` to `tensor_cayley5`, whose inverse has been computed symbolically, and every other -`seq_length`, `0` included, uses the kernel, whatever the length of the data the layer is applied to. +`tensor_cayley2` to `tensor_cayley5`, whose inverse has been computed symbolically. Every other +`seq_length`, `0` included, uses [`tensor_cayley`](@ref), which solves ``(\mathbb{I} + C)\Lambda = +\mathbb{I} - C`` by row-pivoted Gauss–Jordan elimination at any sequence length, on the CPU and on a +device, and in the element type of the data. """ struct VolumePreservingAttention{M, N, ScalarProductType, SL} <: AbstractExplicitLayer{M, N} function VolumePreservingAttention(dim::Int, seq_length::Int = 0; skew_sym::Bool = false) diff --git a/test/devices/jlarray.jl b/test/devices/jlarray.jl index 34c8e8798..6cb059f1d 100644 --- a/test/devices/jlarray.jl +++ b/test/devices/jlarray.jl @@ -7,14 +7,14 @@ using GeometricMachineLearning: _mul, tensor_tensor_mul, tensor_cayley, assign_q using ChainRulesCore: rrule using GPUArraysCore: allowscalar using JLArrays: JLArray, JLBackend -using LinearAlgebra: I import KernelAbstractions using Zygote: gradient using Test import Random -# The bound `10 n eps(T) (1 + ‖Cₖ‖₂²)` of the kernel's orthogonality and value rows. -include("../helpers/cayley_bound.jl") +# The bound `10 n eps(T) (1 + ‖Cₖ‖₂²)` of the kernel's orthogonality and value rows, the dense +# reference `cayley_reference` and the skew batch `skew_batch`. +include("../helpers/cayley.jl") Random.seed!(123) allowscalar(false) @@ -70,8 +70,7 @@ end @testset "tensor_cayley and its pullback on a JLArray" begin for T in (Float32, Float64), n in 2:8 - R = rand(T, n, n, 4) - C = R - permutedims(R, (2, 1, 3)) + C = skew_batch(T, n, 4) seed = rand(T, n, n, 4) D = rand(T, n, n, 4) @@ -79,10 +78,8 @@ end _, C̄ = pullback(JLArray(seed)) @test ondevice(Λ) && ondevice(C̄) - Id = Matrix{T}(I, n, n) - reference = cat(((Id + C[:, :, k]) \ (Id - C[:, :, k]) for k in axes(C, 3))...; dims = 3) - bound = maximum(cayley_bound(T, @view(C[:, :, k]), n) for k in axes(C, 3)) - @test maximum(abs, Array(Λ) - reference) <= bound + bound = cayley_bound(C) + @test maximum(abs, Array(Λ) - cayley_reference(C)) <= bound # The strict rows of check 4, on the device: the `Float64` central difference of the device # pullback with `h = cbrt(eps)`, and the `Float32` pullback against the `Float64` one, which diff --git a/test/devices/metal.jl b/test/devices/metal.jl index 668a1ad2c..df6b2b61d 100644 --- a/test/devices/metal.jl +++ b/test/devices/metal.jl @@ -15,34 +15,14 @@ using GeometricMachineLearning using GeometricMachineLearning: _mul, tensor_tensor_mul, tensor_cayley import KernelAbstractions -using LinearAlgebra: I, dot using Metal using Test import Random -# The bound `10 n eps(T) (1 + ‖Cₖ‖₂²)` of the kernel's orthogonality and value rows. -include("../helpers/cayley_bound.jl") - -# The layer's forward pass on the host, slice by slice: the same scalar products the kernel computes -# for `C`, a dense solve of the Cayley transform, and the product with the data. It returns the -# slices of `C` as well, because the bound of the kernel's orthogonality rows is written in them. -function host_forward(x::AbstractArray{T, 3}, A::AbstractMatrix) where {T} - M = size(x, 1) - out = similar(x) - C = zeros(T, M, M, size(x, 3)) - for k in axes(x, 3) - X = Matrix(view(x, :, :, k)) - L = zeros(T, M, M) - for i in 1:M, j in 1:(i - 1) - - L[i, j] = dot(view(X, :, i), A * view(X, :, j)) - end - C[:, :, k] = (L - L') / T(√M) - Id = Matrix{T}(I, M, M) - out[:, :, k] = X * ((Id + C[:, :, k]) \ (Id - C[:, :, k])) - end - out, C -end +# The bound `10 n eps(T) (1 + ‖Cₖ‖₂²)` of the kernel's orthogonality and value rows, the dense +# reference `cayley_reference`, the skew batch `skew_batch`, and `arbitrary_C`, the host form of the +# layer's `C`. +include("../helpers/cayley.jl") if Metal.functional() Metal.allowscalar(false) @@ -91,25 +71,22 @@ if Metal.functional() # orthogonality rows. @testset "tensor_cayley and VolumePreservingAttention on Metal" begin for n in 2:8 - R = rand(T, n, n, 8) - C = R - permutedims(R, (2, 1, 3)) + C = skew_batch(T, n, 8) Λ = tensor_cayley(MtlArray(C)) @test KernelAbstractions.get_backend(Λ) == backend - Id = Matrix{T}(I, n, n) - reference = cat(((Id + C[:, :, k]) \ (Id - C[:, :, k]) for k in axes(C, 3))...; - dims = 3) - bound = maximum(cayley_bound(T, @view(C[:, :, k]), n) for k in axes(C, 3)) - @test maximum(abs, Array(Λ) - reference) <= bound + @test maximum(abs, Array(Λ) - cayley_reference(C)) <= cayley_bound(C) + # The layer's forward pass against the host: the same `C`, a dense solve of the + # transform, and the product with the data, at the bound scaled by the data. d = VolumePreservingAttention(n, n) ps = NeuralNetwork(Chain(d), CPU(), T).params.L1 x = rand(T, n, n, 8) y = d(MtlArray(x), (A = MtlArray(ps.A),)) - expected, C_layer = host_forward(x, ps.A) - layer_bounds = (cayley_bound(T, @view(C_layer[:, :, k]), n) for k in axes(x, 3)) - layer_bound = maximum(layer_bounds) * maximum(abs, x) + C_layer = arbitrary_C(x, ps.A) + expected = stack(k -> x[:, :, k] * cayley_reference(C_layer[:, :, k]), axes(x, 3)) @test KernelAbstractions.get_backend(y) == backend - @test maximum(abs, Array(y) - expected) <= layer_bound + @test maximum(abs, Array(y) - expected) <= + cayley_bound(C_layer) * maximum(abs, x) end end else diff --git a/test/helpers/cayley.jl b/test/helpers/cayley.jl new file mode 100644 index 000000000..269589610 --- /dev/null +++ b/test/helpers/cayley.jl @@ -0,0 +1,48 @@ +# What the tests of the Cayley kernel `tensor_cayley` and of `VolumePreservingAttention` share: the +# bound the kernel is held to, the dense reference it is compared with, and the inputs the layer +# forms. The kernel's own rows, the device rows and the layer's forward pass all use them. +using LinearAlgebra: I, dot, opnorm + +# The bound `10 n eps(T) (1 + ‖Cₖ‖₂²)` of one slice, and its largest value over a batch. +function cayley_bound(::Type{T}, Ck::AbstractMatrix, n::Int) where {T} + 10.0 * n * eps(T) * (1.0 + Float64(opnorm(Ck, 2))^2) +end + +function cayley_bound(C::AbstractArray{T, 3}) where {T} + maximum(cayley_bound(T, @view(C[:, :, k]), size(C, 1)) for k in axes(C, 3)) +end + +# The Cayley transform `(I + Cₖ) \ (I - Cₖ)` by a dense solve in the element type of `Cₖ`, +# independent of the kernel, of one slice and of every slice of a batch. +function cayley_reference(Ck::AbstractMatrix{T}) where {T} + Id = Matrix{T}(I, size(Ck)) + (Id + Ck) \ (Id - Ck) +end + +function cayley_reference(C::AbstractArray{T, 3}) where {T} + stack(cayley_reference, eachslice(C; dims = 3)) +end + +# A batch of `k` skew `n × n` slices, `R - Rᵀ` for `R = rand(T, n, n, k)`. +function skew_batch(::Type{T}, n::Int, k::Int) where {T} + R = rand(T, n, n, k) + R - permutedims(R, (2, 1, 3)) +end + +# The `:arbitrary` path's `C` on the host, slice by slice: the lower triangle of the scalar products +# of the columns of `x` weighted by `A`, divided by `√M`, and made skew by subtracting its own +# transpose — the same triangle, the same divisor and the same subtraction the layer makes. +function arbitrary_C(x::AbstractArray{T, 3}, A::AbstractMatrix) where {T} + M, seq = size(x, 1), size(x, 2) + C = zeros(T, seq, seq, size(x, 3)) + for k in axes(x, 3) + X = Matrix(@view x[:, :, k]) + L = zeros(T, seq, seq) + for i in 1:seq, j in 1:(i - 1) + + L[i, j] = dot(@view(X[:, i]), A * @view(X[:, j])) + end + C[:, :, k] = (L - L') / T(√M) + end + C +end diff --git a/test/helpers/cayley_bound.jl b/test/helpers/cayley_bound.jl deleted file mode 100644 index 1d08263b2..000000000 --- a/test/helpers/cayley_bound.jl +++ /dev/null @@ -1,8 +0,0 @@ -# The bound of the Cayley kernel's orthogonality and value rows, `10 n eps(T) (1 + ‖Cₖ‖₂²)`, which -# `tensor_cayley`'s docstring states. One definition, because the kernel's own rows, the device rows -# and the layer's forward pass all compare against it and must stay in step with that docstring. -using LinearAlgebra: opnorm - -function cayley_bound(::Type{T}, Ck::AbstractMatrix, n::Int) where {T} - 10.0 * n * eps(T) * (1.0 + Float64(opnorm(Ck, 2))^2) -end diff --git a/test/kernels/cayley.jl b/test/kernels/cayley.jl index 97d37080d..48d4cbd3c 100644 --- a/test/kernels/cayley.jl +++ b/test/kernels/cayley.jl @@ -14,9 +14,9 @@ using LinearAlgebra: BLAS, I, opnorm using Test import Random -# The bound `10 n eps(T) (1 + ‖Cₖ‖₂²)` of the orthogonality and value rows, shared with the device -# rows of the kernel and with the layer's forward pass. -include("../helpers/cayley_bound.jl") +# The bound `10 n eps(T) (1 + ‖Cₖ‖₂²)` of the orthogonality and value rows, and the dense reference +# `cayley_reference`, shared with the device rows of the kernel and with the layer's tests. +include("../helpers/cayley.jl") Random.seed!(1234) @@ -66,18 +66,12 @@ function skew_sym_path(x::AbstractArray{T, 3}, A::AbstractMatrix) where {T} tensor_transpose_tensor_mul(x, _mul(A, x)) end -# The per-slice reference of the orthogonality rows is the solve `(I + Cₖ) \ (I - Cₖ)` in the -# working precision. The value row solves in `BigFloat` and rounds to `T`, so its reference -# approximates the transform itself rather than another elimination in the working precision. -function reference_solve(::Type{T}, Ck::AbstractMatrix) where {T} - Id = Matrix{T}(I, size(Ck)) - (Id + Ck) \ (Id - Ck) -end - +# The per-slice reference of the orthogonality rows is `cayley_reference`, the solve +# `(I + Cₖ) \ (I - Cₖ)` in the working precision. The value row solves in `BigFloat` and rounds to +# `T`, so its reference approximates the transform itself rather than another elimination in the +# working precision. function big_reference_solve(::Type{T}, Ck::AbstractMatrix) where {T} - Id = Matrix{BigFloat}(I, size(Ck)) - C = BigFloat.(Ck) - T.((Id + C) \ (Id - C)) + T.(cayley_reference(BigFloat.(Ck))) end # The orthogonality error of one slice, `‖ΛᵀΛ − I‖∞`: the matrix infinity norm, which is the norm the @@ -94,14 +88,14 @@ function reference_orthogonality_error(C::AbstractArray{T, 3}) where {T} # measured in, so the ratio is a ratio of like quantities. worst = 0.0 for k in axes(C, 3) - worst = max(worst, orthogonality_error(reference_solve(T, @view(C[:, :, k])))) + worst = max(worst, orthogonality_error(cayley_reference(@view(C[:, :, k])))) end worst end -# Every row of the orthogonality check: the ratio against the pivoted reference on all three input -# sets, and the absolute bound on the two exactly skew ones. -function orthogonality_rows(::Type{T}, n::Int, s) where {T} +# Every row of the check, on each of the three input sets: the orthogonality ratio against the +# pivoted reference, the absolute orthogonality bound on the two exactly skew sets, and the value. +function cayley_rows(::Type{T}, n::Int, s) where {T} for (set, C) in (("hand-built", skew_tensor(T, n, s)), ("arbitrary", arbitrary_tensor(T, n, s)), ("skew_sym", skew_sym_tensor(T, n, s))) @@ -124,18 +118,10 @@ function orthogonality_rows(::Type{T}, n::Int, s) where {T} end @test worst <= 1 end - end -end -# The value row: per slice, the distance from the `BigFloat` solve, rounded to `T`, under the same -# bound, on every input set. `Λᵀ` is orthogonal wherever `Λ` is, so the orthogonality rows alone do -# not pin the value. -function value_rows(::Type{T}, n::Int, s) where {T} - for (_, C) in (("hand-built", skew_tensor(T, n, s)), - ("arbitrary", arbitrary_tensor(T, n, s)), - ("skew_sym", skew_sym_tensor(T, n, s))) - @test eltype(C) === T - Λ = tensor_cayley(C) + # The value: per slice, the distance from the `BigFloat` solve, rounded to `T`, under the + # same bound. `Λᵀ` is orthogonal wherever `Λ` is, so the orthogonality rows alone do not pin + # the value. worst = 0.0 for k in axes(C, 3) bound = cayley_bound(T, @view(C[:, :, k]), n) @@ -153,8 +139,8 @@ end x = randn(T, n, n, 7) A = randn(T, n, n) # The layer's own `:arbitrary` operand is exactly skew: one subtraction of a permutation of - # the same array. `M = 4` is the one size here whose `√M` is exact, so the sizes that run - # exercise the rounding of the divisor as well. + # the same array. The sizes leave out `M = 4`, whose `√M` is exact, so that each of them + # exercises the rounding of the divisor as well. C = arbitrary_operand(x, (A = A,)) @test iszero(C + permutedims(C, (2, 1, 3))) # The `:skew_sym` path is skew only to rounding, so nothing is asserted exactly on it. Its @@ -168,8 +154,7 @@ end @testset "tensor_cayley: the orthogonality error and the value, N = $n, s = $s" for T in (Float32, Float64), s in SCALES, n in 2:8 - orthogonality_rows(T, n, s) - value_rows(T, n, s) + cayley_rows(T, n, s) end @testset "tensor_cayley: the pullback, N = $n" for T in (Float32, Float64), n in 2:8 @@ -191,7 +176,7 @@ end # rows: the cotangent is compared relatively, because its scale is the data's. C̄32 = rrule(tensor_cayley, C)[2](Λ̄)[2] C̄64 = rrule(tensor_cayley, Float64.(C))[2](Float64.(Λ̄))[2] - bound = 10 * n * eps(T) * (1 + largest_norm(C)^2) + bound = cayley_bound(C) @test maximum(abs, C̄32 - Float32.(C̄64)) <= bound * maximum(abs, C̄64) end diff --git a/test/kernels/kernel_pullbacks.jl b/test/kernels/kernel_pullbacks.jl index 98d62b0f3..17c75c9d0 100644 --- a/test/kernels/kernel_pullbacks.jl +++ b/test/kernels/kernel_pullbacks.jl @@ -17,11 +17,8 @@ const verbose = false # `n × n` slices that are far from singular, for the inverse rules. invertible(T, n, k) = cat((rand(T, n, n) + n * I for _ in 1:k)...; dims = 3) -# A skew batch for the Cayley rule, whose slices are square. -function skew(T, n, k) - R = rand(T, n, n, k) - R - permutedims(R, (2, 1, 3)) -end +# `skew_batch`, a skew batch for the Cayley rule, whose slices are square. +include("../helpers/cayley.jl") function main(first_dim, second_dim, third_dim, third_tensor_dim) test_rrule(tensor_tensor_mul, rand(first_dim, second_dim, third_tensor_dim), @@ -34,7 +31,7 @@ function main(first_dim, second_dim, third_dim, third_tensor_dim) test_rrule(_mul, rand(first_dim, second_dim)', rand(first_dim, third_dim, third_tensor_dim)) test_rrule(_mul, rand(first_dim, second_dim)', rand(first_dim, third_dim)) test_rrule(init_output, rand(first_dim, first_dim, third_tensor_dim)) - test_rrule(tensor_cayley, skew(Float64, first_dim, third_tensor_dim)) + test_rrule(tensor_cayley, skew_batch(Float64, first_dim, third_tensor_dim)) end const dim_range = 10 @@ -66,7 +63,7 @@ function rules(T) (tensor_inverse4, (invertible(T, 4, 2),)), (tensor_inverse5, (invertible(T, 5, 2),)), (init_output, (rand(T, 3, 3, 2),)), - (tensor_cayley, (skew(T, 3, 2),))) + (tensor_cayley, (skew_batch(T, 3, 2),))) end @testset "a zero seed gives a zero cotangent for every argument" begin diff --git a/test/layers/volume_preserving_attention.jl b/test/layers/volume_preserving_attention.jl index c754be1bc..300b95a1f 100644 --- a/test/layers/volume_preserving_attention.jl +++ b/test/layers/volume_preserving_attention.jl @@ -12,43 +12,20 @@ using FiniteDifferences: central_fdm, j′vp using GPUArraysCore: allowscalar using JLArrays: JLArray, JLBackend import KernelAbstractions -using LinearAlgebra: I, dot using Zygote: gradient using Test import Random # The bound `10 n eps(T) (1 + ‖Cₖ‖₂²)` of the kernel's orthogonality and value rows, which the -# forward pass below is compared against, scaled by the data. -include("../helpers/cayley_bound.jl") +# forward pass below is compared against, scaled by the data; the dense reference +# `cayley_reference`; and `arbitrary_C`, the host form of the `:arbitrary` path's `C`. +include("../helpers/cayley.jl") Random.seed!(1234) allowscalar(false) ondevice(x) = KernelAbstractions.get_backend(x) isa JLBackend -function cayley_reference(C::AbstractMatrix{T}) where {T} - Id = Matrix{T}(I, size(C)) - (Id + C) \ (Id - C) -end - -# The `:arbitrary` path's `C`, slice by slice: the lower triangle of the scalar products of the -# columns of `Z` weighted by `A`, divided by `√M`, and made skew by subtracting its own transpose — -# the same triangle, the same divisor and the same subtraction the layer makes. -function arbitrary_C(x::AbstractArray{T, 3}, A::AbstractMatrix) where {T} - M, seq = size(x, 1), size(x, 2) - C = zeros(T, seq, seq, size(x, 3)) - for k in axes(x, 3) - X = Matrix(@view x[:, :, k]) - L = zeros(T, seq, seq) - for i in 1:seq, j in 1:(i - 1) - - L[i, j] = dot(@view(X[:, i]), A * @view(X[:, j])) - end - C[:, :, k] = (L - L') / T(√M) - end - C -end - function skew_sym_C(x::AbstractArray{T, 3}, A) where {T} C = zeros(T, size(x, 2), size(x, 2), size(x, 3)) D = Matrix(A) @@ -77,12 +54,7 @@ end # The bound of the Cayley kernel's value row, scaled by the largest entry of the data: the product # `Z Λ` multiplies the transform's error by `Z`. -function tolerance(d, x, ps) - T = eltype(x) - C = C_tensor(d, x, ps) - bound = maximum(cayley_bound(T, @view(C[:, :, k]), size(x, 2)) for k in axes(x, 3)) - bound * Float64(maximum(abs, Array(x))) -end +tolerance(d, x, ps) = cayley_bound(C_tensor(d, x, ps)) * Float64(maximum(abs, x)) # A layer of dimension 4 and sequence length `SL` (`0` for the default, which takes data of any # length), its parameters and its data, all on the host; a device case converts them, as @@ -99,10 +71,9 @@ end for N in 2:8, skew_sym in (false, true) d, ps, x = case(N, N, T; skew_sym) - device = VolumePreservingAttention(4, N; skew_sym) ps_device = (A = skew_sym ? SkewSymMatrix(JLArray(ps.A.S), 4) : JLArray(ps.A),) - y_device = device(JLArray(x), ps_device) + y_device = d(JLArray(x), ps_device) @test ondevice(y_device) # The forward pass against the independent host computation. @@ -110,13 +81,13 @@ end @test maximum(abs, Array(y_device) - reference) <= tolerance(d, x, ps) # The gradient with respect to the data, on the device against the host. - gx_device = gradient(z -> sum(device(z, ps_device)), JLArray(x))[1] + gx_device = gradient(z -> sum(d(z, ps_device)), JLArray(x))[1] gx_host = gradient(z -> sum(d(z, ps)), x)[1] @test ondevice(gx_device) @test maximum(abs, Array(gx_device) - gx_host) <= tolerance(d, x, ps) # The gradient with respect to the weight, on the device against the host. - gA_device = gradient(A -> sum(device(JLArray(x), (A = A,))), ps_device.A)[1] + gA_device = gradient(A -> sum(d(JLArray(x), (A = A,))), ps_device.A)[1] gA_host = gradient(A -> sum(d(x, (A = A,))), ps.A)[1] if skew_sym @test ondevice(gA_device.S) diff --git a/test/quality/jet.jl b/test/quality/jet.jl index bfa514642..001b08b8a 100644 --- a/test/quality/jet.jl +++ b/test/quality/jet.jl @@ -66,9 +66,7 @@ const A3{T} = Array{T, 3} zeros(T, 3, 3, 2))) end - # src/kernels/cayley.jl: `tensor_cayley` in test/kernels/cayley.jl, - # test/layers/volume_preserving_attention.jl, test/devices/jlarray.jl and - # test/devices/metal.jl + # src/kernels/cayley.jl: `tensor_cayley`, its kernel and its pullback for T in (Float16, Float32, Float64) @test isempty(JET.get_reports(JET.report_opt(GML.tensor_cayley, (A3{T},); target_modules = TARGET))) From d4416b934d951dee30dc43f7a7660b82faee40ad Mon Sep 17 00:00:00 2001 From: Michael Kraus Date: Sat, 17 Oct 2026 11:50:52 +0200 Subject: [PATCH 10/10] Remove the closed-form Cayley transforms and symbolic inverses for N = 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 --- CHANGELOG.md | 37 +- KNOWN_ISSUES.md | 32 +- scripts/utilities/cayley_timings.jl | 105 ---- scripts/utilities/inverse_generator.jl | 88 ---- src/GeometricMachineLearning.jl | 7 - src/kernels/identity_tensor.jl | 28 - src/kernels/inverses/inverse_2x2.jl | 56 -- src/kernels/inverses/inverse_3x3.jl | 95 ---- src/kernels/inverses/inverse_4x4.jl | 206 -------- src/kernels/inverses/inverse_5x5.jl | 484 ------------------ src/kernels/inverses/tensor_cayley.jl | 19 - .../inverses/tensor_inverse_pullback.jl | 10 - src/layers/volume_preserving_attention.jl | 28 +- .../inverses/test_cayley_transforms.jl | 28 - test/kernels/kernel_pullbacks.jl | 22 +- test/kernels/tensor_cayley.jl | 47 -- test/kernels/tensor_inverse.jl | 123 ----- test/quality/jet.jl | 18 +- test/runtests.jl | 3 - 19 files changed, 9 insertions(+), 1427 deletions(-) delete mode 100644 scripts/utilities/cayley_timings.jl delete mode 100644 scripts/utilities/inverse_generator.jl delete mode 100644 src/kernels/identity_tensor.jl delete mode 100644 src/kernels/inverses/inverse_2x2.jl delete mode 100644 src/kernels/inverses/inverse_3x3.jl delete mode 100644 src/kernels/inverses/inverse_4x4.jl delete mode 100644 src/kernels/inverses/inverse_5x5.jl delete mode 100644 src/kernels/inverses/tensor_cayley.jl delete mode 100644 src/kernels/inverses/tensor_inverse_pullback.jl delete mode 100644 test/kernels/inverses/test_cayley_transforms.jl delete mode 100644 test/kernels/tensor_cayley.jl delete mode 100644 test/kernels/tensor_inverse.jl diff --git a/CHANGELOG.md b/CHANGELOG.md index a8c2291e7..00130a1fd 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -98,47 +98,16 @@ breaking release). ### Changed -- **`VolumePreservingAttention` computes the Cayley transform with the kernel `tensor_cayley` at every sequence length other than 2 to 5, 0 included, and on device arrays; the batch inverse that served those lengths is deleted.** `tensor_cayley(C)` is internal and not exported. It computes `(I + Cₖ)⁻¹(I − Cₖ)` for every slice `k` by row-pivoted Gauss–Jordan elimination, one work item per slice. Its `rrule` computes `C̄ = −(I + Λ)ᵀ Λ̄ (I + Λ)ᵀ / 2` from `Λ` alone, with no second solve and no identity tensor. +- **`VolumePreservingAttention` computes the Cayley transform with the kernel `tensor_cayley` at every sequence length, 0 included, and on device arrays; the batch inverses are deleted.** `tensor_cayley(C)` is internal and not exported. It computes `(I + Cₖ)⁻¹(I − Cₖ)` for every slice `k` by row-pivoted Gauss–Jordan elimination, one work item per slice. Its `rrule` computes `C̄ = −(I + Λ)ᵀ Λ̄ (I + Λ)ᵀ / 2` from `Λ` alone, with no second solve and no identity tensor. - `cpu_inverse`, its `rrule` and `cpu_tensor_cayley` are deleted. The identity tensor that the closed forms use moved from `src/kernels/exponentials/tensor_exponential.jl` to `src/kernels/identity_tensor.jl`, so that the file name matches what it holds; this closes known issue C18, which leaves `KNOWN_ISSUES.md`. + `cpu_inverse`, its `rrule` and `cpu_tensor_cayley` are deleted, and so are the closed forms `tensor_cayley2` to `tensor_cayley5`, the symbolic inverses `tensor_inverse2` to `tensor_inverse5` with their pullback, the identity tensor `init_output`, and `scripts/utilities/inverse_generator.jl`. This closes known issues C18 and K16, which leave `KNOWN_ISSUES.md`. New tests: `test/kernels/cayley.jl` checks the kernel's orthogonality and value, and its pullback against a central difference, for N = 2 to 8. `test/layers/volume_preserving_attention.jl` checks the layer's forward pass and gradients for N = 2 to 8, on the host and on a `JLArray`. `test/devices/jlarray.jl`, `test/devices/metal.jl` and `test/quality/jet.jl` reach the kernel, the last in `Float16`, `Float32` and `Float64`. - `scripts/utilities/cayley_timings.jl` is new. It times `tensor_cayley` against the closed form `tensor_cayleyN` on a batch of 1000 skew slices at scale 1, for N = 2 to 8 and both element types. Each variant, N and element type runs in a fresh process with BLAS on one thread and one untimed warm-up call; a time is the median of five calls, forward and backward. For N = 6 to 8 there is no closed form, so the script prints the kernel's times alone. - - CPU, from two runs of the script on an Apple M-series machine with Julia 1.13.2, at a load average of 3.7 to 4.6. The ratio is the kernel over the closed form. - - | N | T | forward µs, run 1 / run 2 | backward µs, run 1 / run 2 | ratio forward, run 1 / run 2 | ratio backward, run 1 / run 2 | - |:--|:--|--:|--:|--:|--:| - | 2 | Float32 | 23.9 / 24.2 | 31.0 / 31.0 | 0.85 / 0.86 | 0.51 / 0.49 | - | 3 | Float32 | 54.1 / 54.4 | 62.5 / 64.7 | 1.06 / 1.17 | 0.53 / 0.54 | - | 4 | Float32 | 93.2 / 93.6 | 109.2 / 116.0 | 1.24 / 1.31 | 0.54 / 0.57 | - | 5 | Float32 | 139.7 / 140.5 | 173.6 / 169.7 | 1.07 / 1.04 | 0.50 / 0.50 | - | 6 | Float32 | 209.1 / 214.0 | 293.4 / 280.4 | no closed form | no closed form | - | 7 | Float32 | 315.9 / 307.5 | 429.1 / 410.4 | no closed form | no closed form | - | 8 | Float32 | 407.2 / 406.1 | 545.2 / 553.1 | no closed form | no closed form | - | 2 | Float64 | 27.9 / 28.5 | 31.0 / 34.6 | 0.83 / 0.87 | 0.43 / 0.47 | - | 3 | Float64 | 57.4 / 59.2 | 59.9 / 63.3 | 1.25 / 1.42 | 0.40 / 0.41 | - | 4 | Float64 | 94.6 / 93.8 | 111.8 / 119.2 | 1.06 / 0.93 | 0.42 / 0.46 | - | 5 | Float64 | 145.3 / 145.5 | 184.5 / 178.0 | 1.09 / 0.86 | 0.43 / 0.41 | - | 6 | Float64 | 234.8 / 240.6 | 288.8 / 289.8 | no closed form | no closed form | - | 7 | Float64 | 326.3 / 323.6 | 468.4 / 479.0 | no closed form | no closed form | - | 8 | Float64 | 459.7 / 449.1 | 551.0 / 591.6 | no closed form | no closed form | - - At N = 2 the kernel's forward time is below the closed form's in both runs and both element types, - at ratios 0.83 to 0.87. At N = 3 to 5 the forward ratios lie between 0.86 and 1.42, above 1 in some - runs and below 1 in others. The backward ratio is below 1 at N = 2 to 5 in every run. - - Metal, `Float32`, on the `AGXG16CDevice` Metal device, with the script's own `measure` on `MtlArray` - and every timed call synchronized on the device, as the script's header documents. These are - warm-session figures from four runs. The forward ratio at N = 2 is 0.28, 1.28, 0.86 and 0.51, so one - run's ratio moves by up to 4.6× at N = 2. At N = 3 to 5 the forward ratios lie between 0.48 and 1.20. - The backward ratios lie between 0.44 and 0.99 at every N. + The closed forms went because the kernel is not meaningfully slower. Against them, on 1000 skew slices of size N = 2 to 5 on an Apple M-series CPU under load, the kernel's forward time lay between 0.83 and 1.42 of theirs and its backward time between 0.40 and 0.57. On Metal, in `Float32`, the forward ratio lay between 0.28 and 1.28 and the backward ratio between 0.44 and 0.99. These are two cold CPU runs and four warm Metal runs, and one Metal ratio moved by up to 4.6× at N = 2, so the figures are noisy; no ratio exceeded 2. The orthogonality check in `test/kernels/cayley.jl` compares the kernel's `‖ΛₖᵀΛₖ − I‖∞` with that of the pivoted reference `(I + Cₖ) \ (I − Cₖ)`, for N = 2 to 8, both element types and two scales, on three input sets. The ratio of the kernel's error to the reference's is held to a limit of 4 on all three sets. On the two exactly skew sets the kernel's error must also lie within `10·n·eps(T)·(1 + ‖Cₖ‖₂²)`; on the `:skew_sym` set, where `C` is skew only to rounding, the ratio is the only check. The value rows compare each `Λₖ` with a `BigFloat` solve rounded to `T` on the same three sets. - The closed forms `tensor_cayley2` to `tensor_cayley5` and their dispatch are kept for sequence lengths 2 to 5. Every other sequence length, 0 included, uses `tensor_cayley` on the CPU and on devices. The figures for N = 2 to 5 are in the tables above. - - `Pkg.test()` no longer runs the doctests. `test/quality/doctests.jl` is now the group `doctests`, which an empty `ARGS` does not run; `Pkg.test(test_args = ["doctests"])` runs it. In CI the Doctests job stays their runner, so the test matrix no longer runs them a second time. - **CI uploads coverage from the `Julia 1 - ubuntu-latest` job instead of `Julia min`, and a test job saves the Julia cache only when it succeeds.** - **The package needs Julia 1.12 or later** (`[compat] julia = "1.12"`, was `"1.11"`). On Julia 1.11.9, the gradient of an `LNNLoss` with respect to the network parameters stops with `StackOverflowError` in type inference ("Internal error: during type inference of … Encountered stack overflow") on macOS and Windows; the `Julia min` job on Linux passes. The overflow comes from SymbolicUtils 4.49.0, whose `rrule` for `SymbolicUtils.Code.create_array` returns an exactly typed tuple of tangents, one per array element; Zygote's reverse pass through the generated Hessian of `lagrangian_acceleration` is then inferred far deeper. With the 4.48.0 rule restored, or in a task with a 64 MB stack, the same gradient passes on 1.11.9; Julia 1.12.7 and 1.13.1 pass it on the root task. This closes #327: the `Julia min` jobs on macOS and Windows had failed in the *Lagrangian Neural Network* testset since 2026-10-01, and every testset after it was never run there. diff --git a/KNOWN_ISSUES.md b/KNOWN_ISSUES.md index 85bf93886..5ff61182a 100644 --- a/KNOWN_ISSUES.md +++ b/KNOWN_ISSUES.md @@ -198,32 +198,6 @@ What is known to be wrong in GeometricMachineLearning and is not fixed. belongs with the data signatures of the losses as a whole: either each loss accepts only the data its vector field can take, or `(q, p)` data is converted before it reaches the field. -### K16 · `scripts/utilities/inverse_generator.jl` reproduces the wrappers of the committed inverse files byte for byte, but not their kernels. - -- **location:** `scripts/utilities/inverse_generator.jl:1` -- **kind:** found late -- **found:** 2026-10-15, by the builder of part M3; the same output comes from the generator of `origin/main` at `2a63f26a` -- **evidence:** - Formatted with `style = "sciml"`, the generator writes the kernel of each of the four - `inverse_NxN.jl` files, the part before `function tensor_inverseN(`, with the factors of some - products in another order than the committed file, from line 6 or 7 on: - - ``` - < t2 = (*)((*)(t1, (getindex)(A, 2, 1, k)), (getindex)(A, 1, 2, k)) - --- - > t2 = (*)((*)(t1, (getindex)(A, 1, 2, k)), (getindex)(A, 2, 1, k)) - ``` - - The order depends on the Julia version, not on the package versions. With Symbolics 7.39.2, - SymbolicUtils 4.46.6 and JuliaFormatter 2.14.0 fixed, Julia 1.11.9, 1.12.6, 1.12.7 and 1.13.1 - each give a different order, and none of them reproduces all four committed kernels; 1.12.6 - reproduces the 2×2 one. On Julia 1.13.1, SymbolicUtils 4.46.6 and 4.51.0 give identical output, - and two runs on one Julia version give identical output. The Julia build that made the committed - kernels on 2026-09-14 is no longer installed, so that run cannot be repeated. - - The rest of each file, from `function tensor_inverseN(` to the end, is byte for byte what the - generator's `wrapper_source(n)` writes, and the script's header claims no more than that. - ### K17 · `assign_q_and_p` and `PoissonTensor * v` drop the rows from `2N + 1` on without an error. - **location:** `src/kernels/assign_q_and_p.jl` @@ -245,8 +219,7 @@ What is known to be wrong in GeometricMachineLearning and is not fixed. - **evidence:** On 200 skew `Float32` slices of `N = 5` whose largest `‖Cₖ‖₂` is 6.55e8, `tensor_cayley` returns non-finite values in 23 of them and raises no error. The scale is far outside check 3's `{1, 1e2}`, - and no method is accurate there: `tensor_cayley5`, the closed form that stays, is non-finite in 197 - of the same 200 slices, and the dense solve `(I + Cₖ) \ (I − Cₖ)` is finite on all 200 but has an + and no method is accurate there: the dense solve `(I + Cₖ) \ (I − Cₖ)` is finite on all 200 but has an orthogonality error `‖ΛₖᵀΛₖ − I‖∞` of 7.3e4 at worst, which is no answer at that scale. The cause is cancellation in the last pivot, not overflow: in all 23 non-finite slices the last `Float32` pivot is exactly `0`. The same elimination in `Float64` gives a last pivot of magnitude 1.7 to 15 @@ -256,14 +229,13 @@ What is known to be wrong in GeometricMachineLearning and is not fixed. slices is at most 284. ```julia - using GeometricMachineLearning: tensor_cayley, tensor_cayley5 + using GeometricMachineLearning: tensor_cayley import Random Random.seed!(3) R = 1.0f8 * randn(Float32, 5, 5, 200) C = R - permutedims(R, (2, 1, 3)) nonfinite(L) = count(k -> !all(isfinite, L[:, :, k]), 1:200) nonfinite(tensor_cayley(C)) # 23 of 200 - nonfinite(tensor_cayley5(C)) # 197 of 200 ``` ## Upstream diff --git a/scripts/utilities/cayley_timings.jl b/scripts/utilities/cayley_timings.jl deleted file mode 100644 index ba84d7192..000000000 --- a/scripts/utilities/cayley_timings.jl +++ /dev/null @@ -1,105 +0,0 @@ -# The forward and the backward time of the Cayley transform of a batch of 1000 skew slices at scale -# 1, `tensor_cayley`, against the closed form the tree defines for that slice size: `tensor_cayleyN` -# for N = 2 to 5. A size whose closed form the tree does not define, N = 6 to 8, prints the kernel's -# own time and no ratio. BLAS is pinned to one thread, and every variant, every N and every element -# type runs in a fresh process; a process makes one warm-up call it does not time, then takes the -# median of five timed calls. The parent prints the ratio of the kernel to the closed form, forward -# and backward. -# -# Run from the repository root, with the package's own environment, which carries `Zygote` and the -# standard libraries this script uses; `--project=scripts` resolves them as well, and CUDA with them: -# julia --startup-file=no --project=. scripts/utilities/cayley_timings.jl -# julia --startup-file=no --project=. scripts/utilities/cayley_timings.jl tensor_cayley 4 Float32 -# -# Metal rows for N = 6 to 8 print the kernel's time alone. -# N = 2 to 5 compare against the closed forms, which run on Metal as well; `Metal` is not a -# dependency of the package, so its column runs `measure` in a process whose environment carries -# Metal, which `scripts/Project.toml` does not. This file runs its table only as the program, so a -# session reaches `measure` by including it and calling it with the constructor that moves the batch -# and the seed onto the device and with the wrapper that makes a timed call a time rather than a -# launch: -# -# julia> include("scripts/utilities/cayley_timings.jl") -# julia> sync(x) = (Metal.synchronize(); x) -# julia> measure("tensor_cayley", 6, Float32, sync, MtlArray) -# -# `Metal.synchronize` is what makes the figure a time rather than a launch: without it the call -# returns before the kernel has run. - -using GeometricMachineLearning -using LinearAlgebra: BLAS -using Printf -import Random -import Zygote - -const KERNEL = "tensor_cayley" -const SLICES = 1000 -const RUNS = 5 - -median(v) = sort(v)[(length(v) + 1) ÷ 2] - -# One untimed warm-up call, then `RUNS` timed calls, in µs. `sync` wraps one timed call and is the -# identity except on a device, where the figure is a time rather than a launch. -function times(f, sync = identity) - sync(f()) - median([(@elapsed sync(f())) * 1e6 for _ in 1:RUNS]) -end - -function skew_batch(::Type{T}, n) where {T} - R = randn(T, n, n, SLICES) - R - permutedims(R, (2, 1, 3)) -end - -# The forward and the backward time of one variant in this process, in µs. `array` moves the batch -# and the seed to the backend the variant runs on, and is `Array` on the CPU. -function measure(variant, n, ::Type{T}, sync = identity, array = Array) where {T} - BLAS.set_num_threads(1) - Random.seed!(0) - - C = array(skew_batch(T, n)) - f = getfield(GeometricMachineLearning, Symbol(variant)) - - forward = times(() -> f(C), sync) - Λ, pullback = Zygote.pullback(f, C) - Λ̄ = array(rand(T, size(Λ))) - backward = times(() -> pullback(Λ̄), sync) - - forward, backward -end - -# The closed form of the transform of an `n × n` slice, where the tree defines one. -reference(n) = n <= 5 ? "tensor_cayley$n" : nothing - -function one_process(variant, n, T) - cmd = `$(Base.julia_cmd()) --startup-file=no --project=$(Base.active_project()) $(@__FILE__) $variant $n $T` - parse.(Float64, split(readchomp(cmd))) -end - -function table() - @printf("%-17s %-2s %-8s %11s %11s %s\n", - "variant", "N", "T", "forward µs", "backward µs", - "ratio kernel/closed form") - for T in (Float32, Float64), n in 2:8 - - kernel = one_process(KERNEL, n, T) - name = reference(n) - if name !== nothing && isdefined(GeometricMachineLearning, Symbol(name)) - closed = one_process(name, n, T) - @printf("%-17s %-2d %-8s %11.1f %11.1f forward %.2f, backward %.2f\n", - KERNEL, n, T, kernel[1], kernel[2], kernel[1] / closed[1], - kernel[2] / closed[2]) - else - @printf("%-17s %-2d %-8s %11.1f %11.1f no closed form\n", - KERNEL, n, T, kernel[1], kernel[2]) - end - end -end - -if abspath(PROGRAM_FILE) == abspath(@__FILE__) - if length(ARGS) == 3 - variant, n, T = ARGS[1], parse(Int, ARGS[2]), getfield(Base, Symbol(ARGS[3])) - println(join(measure(variant, n, T), " ")) - else - table() - end -end diff --git a/scripts/utilities/inverse_generator.jl b/scripts/utilities/inverse_generator.jl deleted file mode 100644 index 809d30b32..000000000 --- a/scripts/utilities/inverse_generator.jl +++ /dev/null @@ -1,88 +0,0 @@ -# Generates the four `src/kernels/inverses/inverse_NxN.jl` files. The script writes them into the -# working directory; run JuliaFormatter over the output with `style = "sciml"` and copy it into -# place. The part of each committed file from `function tensor_inverseN(` to the end is what -# `wrapper_source(n)` writes, byte for byte. The kernel above it is not reproduced byte for byte: -# the order of the factors in its products changes with the Julia version (`KNOWN_ISSUES.md`, K16). -# -# `Symbolics` is deliberately not a dependency of the package: the generated kernels are committed -# so that building the package never needs a symbolic stack. Run the script by hand, in an -# environment that has `Symbolics` and `JuliaFormatter`. -# -# `cse = true` is load-bearing. `build_function` defaults it to `false`, and each entry of the -# inverse then comes out as one deeply nested expression. JuliaFormatter indents every level of -# that nesting, and the cost is quadratic in the depth: formatted without CSE, the 5x5 file is -# 25 MB, of which 98 % is leading whitespace, and `fatou lint` does not terminate on it. With CSE -# the same inverse is a few hundred straight-line assignments and formats to tens of kilobytes. - -using Symbolics - -function matrix_inverse(n) - @variables A[1:n, 1:n] - - B = inv(collect(A)) - - # `simplify.(B)` is deliberately absent. Under Symbolics 7.39.2 it changes nothing here: the - # generated string is byte-identical with and without it at every n this script emits. It costs - # 0.7 s at n = 5, and it becomes impractical above n = 6. - - build_function(B, A; cse = true)[2] -end - -# `build_function` emits a plain `function (ˍ₋out, A)` over one n×n matrix. The kernels run over a -# tensor, one matrix per slice `k`, so every read of `A` gains the slice index and the linear -# output index becomes a Cartesian one. -function kernel_source(n) - expr = matrix_inverse(n) - str = replace(replace(replace(string(expr), r"#= [^*\s]* =#" => ""), r"\n[\s]*\n" => "\n"), "Num" => "") - str = replace(str, r"var\"##cse#(\d+)\"" => s"t\1") - str = replace(str, r"\(getindex\)\(A, (\d+), (\d+)\)" => s"(getindex)(A, \1, \2, k)") - str = replace(str, r"ˍ₋out\[(\d+)\]" => - m -> begin - l = parse(Int, match(r"\d+", m).match) - "ˍ₋out[$(mod1(l, n)), $(cld(l, n)), k]" - end) - str = replace(str, r"^(\s*)ˍ₋out$"m => s"\1nothing") - replace(str, - "function (ˍ₋out, A)" => "@kernel function inv$(n)$(n)_kernel!(ˍ₋out, A)\n k = @index(Global)") -end - -function wrapper_source(n) - """ - function tensor_inverse$(n)(A::AbstractArray{T, 3}) where {T} - out = similar(A) - - tensor_inverse$(n)!(out, A) - - out - end - - function tensor_inverse$(n)!(out::AbstractArray{T, 3}, A::AbstractArray{T, 3}) where {T} - @assert size(A, 1) == size(A, 2) == $(n) - @assert size(A) == size(out) - - backend = networkbackend(out) - inv$(n)$(n)! = inv$(n)$(n)_kernel!(backend) - - inv$(n)$(n)!(out, A, ndrange = size(A, 3)) - - nothing - end - - function ChainRulesCore.rrule(::typeof(tensor_inverse$(n)), A::AT) where { - T, AT <: AbstractArray{T, 3}} - out = tensor_inverse$(n)(A) - - function tensor_inverse_pullback(out_diff) - _tensor_inverse_pullback(out, out_diff) - end - out, tensor_inverse_pullback - end - """ -end - -# The four sizes the package ships. `VolumePreservingAttention` dispatches on the sequence length -# and falls back to `tensor_cayley` for every other one. -for n in 2:5 - println("$(n)x$(n)...") - write("inverse_$(n)x$(n).jl", kernel_source(n) * "\n\n" * wrapper_source(n)) -end diff --git a/src/GeometricMachineLearning.jl b/src/GeometricMachineLearning.jl index 379913c8d..2758b4f70 100644 --- a/src/GeometricMachineLearning.jl +++ b/src/GeometricMachineLearning.jl @@ -128,13 +128,6 @@ include("kernels/tensor_transpose_tensor_mul.jl") include("kernels/tensor_tensor_transpose_mul.jl") include("kernels/mul.jl") include("kernels/cayley.jl") -include("kernels/identity_tensor.jl") -include("kernels/inverses/tensor_inverse_pullback.jl") -include("kernels/inverses/inverse_2x2.jl") -include("kernels/inverses/inverse_3x3.jl") -include("kernels/inverses/inverse_4x4.jl") -include("kernels/inverses/inverse_5x5.jl") -include("kernels/inverses/tensor_cayley.jl") include("kernels/inverses/tensor_mat_skew_sym_assign.jl") include("kernels/kernel_ad_routines/seed.jl") diff --git a/src/kernels/identity_tensor.jl b/src/kernels/identity_tensor.jl deleted file mode 100644 index b53fcdbf2..000000000 --- a/src/kernels/identity_tensor.jl +++ /dev/null @@ -1,28 +0,0 @@ -# The identity tensor and its `rrule`, which the closed forms of the Cayley transform of a batch of -# small slices, `tensor_cayley2` to `tensor_cayley5`, use to form `I ± C`. - -function init_output(B::AbstractArray{T, 3}) where {T} - output = zero(B) - assign_ones!(output) - output -end - -function assign_ones!(output::AbstractArray{T, 3}) where {T} - backend = networkbackend(output) - assign_ones_backend! = assign_ones_kernel!(backend) - dims = (size(output, 1), size(output, 3)) - assign_ones_backend!(output, ndrange = dims) -end - -@kernel function assign_ones_kernel!(output::AbstractArray{T, 3}) where {T} - i, k = @index(Global, NTuple) - output[i, i, k] = one(T) -end - -function ChainRulesCore.rrule(::typeof(init_output), B::AbstractArray{T, 3}) where {T} - output = init_output(B) - function init_output_pullback(_output_diff) - return NoTangent(), ZeroTangent() - end - output, init_output_pullback -end diff --git a/src/kernels/inverses/inverse_2x2.jl b/src/kernels/inverses/inverse_2x2.jl deleted file mode 100644 index 0c80ecdcf..000000000 --- a/src/kernels/inverses/inverse_2x2.jl +++ /dev/null @@ -1,56 +0,0 @@ -@kernel function inv22_kernel!(ˍ₋out, A) - k = @index(Global) - @inbounds begin - begin - begin - t1 = -1 - t2 = (*)((*)(t1, (getindex)(A, 1, 2, k)), (getindex)(A, 2, 1, k)) - t3 = (*)((getindex)(A, 1, 1, k), (getindex)(A, 2, 2, k)) - t4 = (+)(t2, t3) - t5 = (/)((getindex)(A, 2, 2, k), t4) - t6 = (*)(t1, (getindex)(A, 2, 1, k)) - t7 = (/)(t6, t4) - t8 = (*)(t1, (getindex)(A, 1, 2, k)) - t9 = (/)(t8, t4) - t10 = (/)((getindex)(A, 1, 1, k), t4) - @inbounds begin - ˍ₋out[1, 1, k] = t5 - ˍ₋out[2, 1, k] = t7 - ˍ₋out[1, 2, k] = t9 - ˍ₋out[2, 2, k] = t10 - nothing - end - end - end - end -end - -function tensor_inverse2(A::AbstractArray{T, 3}) where {T} - out = similar(A) - - tensor_inverse2!(out, A) - - out -end - -function tensor_inverse2!(out::AbstractArray{T, 3}, A::AbstractArray{T, 3}) where {T} - @assert size(A, 1) == size(A, 2) == 2 - @assert size(A) == size(out) - - backend = networkbackend(out) - inv22! = inv22_kernel!(backend) - - inv22!(out, A, ndrange = size(A, 3)) - - nothing -end - -function ChainRulesCore.rrule(::typeof(tensor_inverse2), A::AT) where { - T, AT <: AbstractArray{T, 3}} - out = tensor_inverse2(A) - - function tensor_inverse_pullback(out_diff) - _tensor_inverse_pullback(out, out_diff) - end - out, tensor_inverse_pullback -end diff --git a/src/kernels/inverses/inverse_3x3.jl b/src/kernels/inverses/inverse_3x3.jl deleted file mode 100644 index f8afcfad7..000000000 --- a/src/kernels/inverses/inverse_3x3.jl +++ /dev/null @@ -1,95 +0,0 @@ -@kernel function inv33_kernel!(ˍ₋out, A) - k = @index(Global) - @inbounds begin - begin - begin - t1 = -1 - t2 = (*)((*)(t1, (getindex)(A, 2, 3, k)), (getindex)(A, 3, 2, k)) - t3 = (*)((getindex)(A, 3, 3, k), (getindex)(A, 2, 2, k)) - t4 = (+)(t2, t3) - t5 = (*)((*)(t1, (getindex)(A, 2, 3, k)), (getindex)(A, 3, 1, k)) - t6 = (*)((getindex)(A, 3, 3, k), (getindex)(A, 2, 1, k)) - t7 = (+)(t5, t6) - t8 = (*)((*)(t1, t7), (getindex)(A, 1, 2, k)) - t9 = (*)((getindex)(A, 1, 1, k), t4) - t10 = (*)((getindex)(A, 2, 1, k), (getindex)(A, 3, 2, k)) - t11 = (*)((*)(t1, (getindex)(A, 3, 1, k)), (getindex)(A, 2, 2, k)) - t12 = (+)(t10, t11) - t13 = (*)(t12, (getindex)(A, 1, 3, k)) - t14 = (+)((+)(t8, t9), t13) - t15 = (/)(t4, t14) - t16 = (*)((getindex)(A, 2, 3, k), (getindex)(A, 3, 1, k)) - t17 = (*)((*)(t1, (getindex)(A, 3, 3, k)), (getindex)(A, 2, 1, k)) - t18 = (+)(t16, t17) - t19 = (/)(t18, t14) - t20 = (/)(t12, t14) - t21 = (*)((getindex)(A, 3, 2, k), (getindex)(A, 1, 3, k)) - t22 = (*)((*)(t1, (getindex)(A, 1, 2, k)), (getindex)(A, 3, 3, k)) - t23 = (+)(t21, t22) - t24 = (/)(t23, t14) - t25 = (*)((getindex)(A, 3, 3, k), (getindex)(A, 1, 1, k)) - t26 = (*)((*)(t1, (getindex)(A, 3, 1, k)), (getindex)(A, 1, 3, k)) - t27 = (+)(t25, t26) - t28 = (/)(t27, t14) - t29 = (*)((getindex)(A, 1, 2, k), (getindex)(A, 3, 1, k)) - t30 = (*)((*)(t1, (getindex)(A, 1, 1, k)), (getindex)(A, 3, 2, k)) - t31 = (+)(t29, t30) - t32 = (/)(t31, t14) - t33 = (*)((getindex)(A, 2, 3, k), (getindex)(A, 1, 2, k)) - t34 = (*)((*)(t1, (getindex)(A, 2, 2, k)), (getindex)(A, 1, 3, k)) - t35 = (+)(t33, t34) - t36 = (/)(t35, t14) - t37 = (*)((getindex)(A, 2, 1, k), (getindex)(A, 1, 3, k)) - t38 = (*)((*)(t1, (getindex)(A, 2, 3, k)), (getindex)(A, 1, 1, k)) - t39 = (+)(t37, t38) - t40 = (/)(t39, t14) - t41 = (*)((*)(t1, (getindex)(A, 1, 2, k)), (getindex)(A, 2, 1, k)) - t42 = (*)((getindex)(A, 1, 1, k), (getindex)(A, 2, 2, k)) - t43 = (+)(t41, t42) - t44 = (/)(t43, t14) - @inbounds begin - ˍ₋out[1, 1, k] = t15 - ˍ₋out[2, 1, k] = t19 - ˍ₋out[3, 1, k] = t20 - ˍ₋out[1, 2, k] = t24 - ˍ₋out[2, 2, k] = t28 - ˍ₋out[3, 2, k] = t32 - ˍ₋out[1, 3, k] = t36 - ˍ₋out[2, 3, k] = t40 - ˍ₋out[3, 3, k] = t44 - nothing - end - end - end - end -end - -function tensor_inverse3(A::AbstractArray{T, 3}) where {T} - out = similar(A) - - tensor_inverse3!(out, A) - - out -end - -function tensor_inverse3!(out::AbstractArray{T, 3}, A::AbstractArray{T, 3}) where {T} - @assert size(A, 1) == size(A, 2) == 3 - @assert size(A) == size(out) - - backend = networkbackend(out) - inv33! = inv33_kernel!(backend) - - inv33!(out, A, ndrange = size(A, 3)) - - nothing -end - -function ChainRulesCore.rrule(::typeof(tensor_inverse3), A::AT) where { - T, AT <: AbstractArray{T, 3}} - out = tensor_inverse3(A) - - function tensor_inverse_pullback(out_diff) - _tensor_inverse_pullback(out, out_diff) - end - out, tensor_inverse_pullback -end diff --git a/src/kernels/inverses/inverse_4x4.jl b/src/kernels/inverses/inverse_4x4.jl deleted file mode 100644 index a52c72a83..000000000 --- a/src/kernels/inverses/inverse_4x4.jl +++ /dev/null @@ -1,206 +0,0 @@ -@kernel function inv44_kernel!(ˍ₋out, A) - k = @index(Global) - @inbounds begin - begin - begin - t1 = (*)((getindex)(A, 4, 4, k), (getindex)(A, 3, 3, k)) - t2 = -1 - t3 = (*)((*)(t2, (getindex)(A, 4, 3, k)), (getindex)(A, 3, 4, k)) - t4 = (+)(t1, t3) - t5 = (*)((getindex)(A, 2, 2, k), t4) - t6 = (*)((getindex)(A, 4, 4, k), (getindex)(A, 3, 2, k)) - t7 = (*)((*)(t2, (getindex)(A, 4, 2, k)), (getindex)(A, 3, 4, k)) - t8 = (+)(t6, t7) - t9 = (*)((*)(t2, t8), (getindex)(A, 2, 3, k)) - t10 = (*)((*)(t2, (getindex)(A, 4, 2, k)), (getindex)(A, 3, 3, k)) - t11 = (*)((getindex)(A, 4, 3, k), (getindex)(A, 3, 2, k)) - t12 = (+)(t10, t11) - t13 = (*)(t12, (getindex)(A, 2, 4, k)) - t14 = (+)((+)(t5, t9), t13) - t15 = (*)((getindex)(A, 2, 1, k), t4) - t16 = (*)((*)(t2, (getindex)(A, 4, 1, k)), (getindex)(A, 3, 3, k)) - t17 = (*)((getindex)(A, 4, 3, k), (getindex)(A, 3, 1, k)) - t18 = (+)(t16, t17) - t19 = (*)(t18, (getindex)(A, 2, 4, k)) - t20 = (*)((*)(t2, (getindex)(A, 4, 1, k)), (getindex)(A, 3, 4, k)) - t21 = (*)((getindex)(A, 3, 1, k), (getindex)(A, 4, 4, k)) - t22 = (+)(t20, t21) - t23 = (*)((*)(t2, t22), (getindex)(A, 2, 3, k)) - t24 = (+)((+)(t15, t19), t23) - t25 = (*)((*)(t2, (getindex)(A, 1, 2, k)), t24) - t26 = (*)(t14, (getindex)(A, 1, 1, k)) - t27 = (*)((*)(t2, (getindex)(A, 4, 1, k)), (getindex)(A, 3, 2, k)) - t28 = (*)((getindex)(A, 4, 2, k), (getindex)(A, 3, 1, k)) - t29 = (+)(t27, t28) - t30 = (*)((getindex)(A, 2, 3, k), t29) - t31 = (*)((getindex)(A, 2, 1, k), t12) - t32 = (*)((*)(t2, (getindex)(A, 2, 2, k)), t18) - t33 = (+)((+)(t30, t31), t32) - t34 = (*)((*)(t2, (getindex)(A, 1, 4, k)), t33) - t35 = (*)((*)(t2, (getindex)(A, 2, 2, k)), t22) - t36 = (*)((getindex)(A, 2, 1, k), t8) - t37 = (*)((getindex)(A, 2, 4, k), t29) - t38 = (+)((+)(t35, t36), t37) - t39 = (*)(t38, (getindex)(A, 1, 3, k)) - t40 = (+)((+)((+)(t25, t26), t34), t39) - t41 = (/)(t14, t40) - t42 = (*)((*)(t2, (getindex)(A, 2, 1, k)), t4) - t43 = (*)((*)(t2, t18), (getindex)(A, 2, 4, k)) - t44 = (*)(t22, (getindex)(A, 2, 3, k)) - t45 = (+)((+)(t42, t43), t44) - t46 = (/)(t45, t40) - t47 = (/)(t38, t40) - t48 = (*)((*)(t2, (getindex)(A, 2, 3, k)), t29) - t49 = (*)((*)(t2, (getindex)(A, 2, 1, k)), t12) - t50 = (*)((getindex)(A, 2, 2, k), t18) - t51 = (+)((+)(t48, t49), t50) - t52 = (/)(t51, t40) - t53 = (*)(t8, (getindex)(A, 1, 3, k)) - t54 = (*)((*)(t2, (getindex)(A, 1, 2, k)), t4) - t55 = (*)((*)(t2, (getindex)(A, 1, 4, k)), t12) - t56 = (+)((+)(t53, t54), t55) - t57 = (/)(t56, t40) - t58 = (*)(t18, (getindex)(A, 1, 4, k)) - t59 = (*)(t4, (getindex)(A, 1, 1, k)) - t60 = (*)((*)(t2, t22), (getindex)(A, 1, 3, k)) - t61 = (+)((+)(t58, t59), t60) - t62 = (/)(t61, t40) - t63 = (*)((*)(t2, t8), (getindex)(A, 1, 1, k)) - t64 = (*)((*)(t2, (getindex)(A, 1, 4, k)), t29) - t65 = (*)((getindex)(A, 1, 2, k), t22) - t66 = (+)((+)(t63, t64), t65) - t67 = (/)(t66, t40) - t68 = (*)((*)(t2, (getindex)(A, 1, 2, k)), t18) - t69 = (*)(t29, (getindex)(A, 1, 3, k)) - t70 = (*)(t12, (getindex)(A, 1, 1, k)) - t71 = (+)((+)(t68, t69), t70) - t72 = (/)(t71, t40) - t73 = (*)((*)(t2, (getindex)(A, 4, 3, k)), (getindex)(A, 2, 4, k)) - t74 = (*)((getindex)(A, 4, 4, k), (getindex)(A, 2, 3, k)) - t75 = (+)(t73, t74) - t76 = (*)((getindex)(A, 1, 2, k), t75) - t77 = (*)((*)(t2, (getindex)(A, 4, 2, k)), (getindex)(A, 2, 4, k)) - t78 = (*)((getindex)(A, 2, 2, k), (getindex)(A, 4, 4, k)) - t79 = (+)(t77, t78) - t80 = (*)((*)(t2, t79), (getindex)(A, 1, 3, k)) - t81 = (*)((*)(t2, (getindex)(A, 4, 2, k)), (getindex)(A, 2, 3, k)) - t82 = (*)((getindex)(A, 4, 3, k), (getindex)(A, 2, 2, k)) - t83 = (+)(t81, t82) - t84 = (*)((getindex)(A, 1, 4, k), t83) - t85 = (+)((+)(t76, t80), t84) - t86 = (/)(t85, t40) - t87 = (*)((getindex)(A, 2, 1, k), (getindex)(A, 4, 4, k)) - t88 = (*)((*)(t2, (getindex)(A, 4, 1, k)), (getindex)(A, 2, 4, k)) - t89 = (+)(t87, t88) - t90 = (*)(t89, (getindex)(A, 1, 3, k)) - t91 = (*)((getindex)(A, 4, 3, k), (getindex)(A, 2, 1, k)) - t92 = (*)((*)(t2, (getindex)(A, 4, 1, k)), (getindex)(A, 2, 3, k)) - t93 = (+)(t91, t92) - t94 = (*)((*)(t2, (getindex)(A, 1, 4, k)), t93) - t95 = (*)((*)(t2, t75), (getindex)(A, 1, 1, k)) - t96 = (+)((+)(t90, t94), t95) - t97 = (/)(t96, t40) - t98 = (*)(t79, (getindex)(A, 1, 1, k)) - t99 = (*)((*)(t2, t89), (getindex)(A, 1, 2, k)) - t100 = (*)((*)(t2, (getindex)(A, 4, 1, k)), (getindex)(A, 2, 2, k)) - t101 = (*)((getindex)(A, 2, 1, k), (getindex)(A, 4, 2, k)) - t102 = (+)(t100, t101) - t103 = (*)(t102, (getindex)(A, 1, 4, k)) - t104 = (+)((+)(t98, t99), t103) - t105 = (/)(t104, t40) - t106 = (*)((getindex)(A, 1, 2, k), t93) - t107 = (*)((*)(t2, t102), (getindex)(A, 1, 3, k)) - t108 = (*)((*)(t2, t83), (getindex)(A, 1, 1, k)) - t109 = (+)((+)(t106, t107), t108) - t110 = (/)(t109, t40) - t111 = (*)((getindex)(A, 3, 4, k), (getindex)(A, 2, 2, k)) - t112 = (*)((*)(t2, (getindex)(A, 2, 4, k)), (getindex)(A, 3, 2, k)) - t113 = (+)(t111, t112) - t114 = (*)(t113, (getindex)(A, 1, 3, k)) - t115 = (*)((*)(t2, (getindex)(A, 2, 3, k)), (getindex)(A, 3, 2, k)) - t116 = (*)((getindex)(A, 2, 2, k), (getindex)(A, 3, 3, k)) - t117 = (+)(t115, t116) - t118 = (*)((*)(t2, t117), (getindex)(A, 1, 4, k)) - t119 = (*)((*)(t2, (getindex)(A, 2, 4, k)), (getindex)(A, 3, 3, k)) - t120 = (*)((getindex)(A, 3, 4, k), (getindex)(A, 2, 3, k)) - t121 = (+)(t119, t120) - t122 = (*)((*)(t2, (getindex)(A, 1, 2, k)), t121) - t123 = (+)((+)(t114, t118), t122) - t124 = (/)(t123, t40) - t125 = (*)((getindex)(A, 2, 1, k), (getindex)(A, 3, 4, k)) - t126 = (*)((*)(t2, (getindex)(A, 3, 1, k)), (getindex)(A, 2, 4, k)) - t127 = (+)(t125, t126) - t128 = (*)((*)(t2, t127), (getindex)(A, 1, 3, k)) - t129 = (*)(t121, (getindex)(A, 1, 1, k)) - t130 = (*)((getindex)(A, 2, 1, k), (getindex)(A, 3, 3, k)) - t131 = (*)((*)(t2, (getindex)(A, 3, 1, k)), (getindex)(A, 2, 3, k)) - t132 = (+)(t130, t131) - t133 = (*)((getindex)(A, 1, 4, k), t132) - t134 = (+)((+)(t128, t129), t133) - t135 = (/)(t134, t40) - t136 = (*)((getindex)(A, 1, 2, k), t127) - t137 = (*)((getindex)(A, 2, 1, k), (getindex)(A, 3, 2, k)) - t138 = (*)((*)(t2, (getindex)(A, 3, 1, k)), (getindex)(A, 2, 2, k)) - t139 = (+)(t137, t138) - t140 = (*)((*)(t2, (getindex)(A, 1, 4, k)), t139) - t141 = (*)((*)(t2, t113), (getindex)(A, 1, 1, k)) - t142 = (+)((+)(t136, t140), t141) - t143 = (/)(t142, t40) - t144 = (*)(t139, (getindex)(A, 1, 3, k)) - t145 = (*)((*)(t2, (getindex)(A, 1, 2, k)), t132) - t146 = (*)(t117, (getindex)(A, 1, 1, k)) - t147 = (+)((+)(t144, t145), t146) - t148 = (/)(t147, t40) - @inbounds begin - ˍ₋out[1, 1, k] = t41 - ˍ₋out[2, 1, k] = t46 - ˍ₋out[3, 1, k] = t47 - ˍ₋out[4, 1, k] = t52 - ˍ₋out[1, 2, k] = t57 - ˍ₋out[2, 2, k] = t62 - ˍ₋out[3, 2, k] = t67 - ˍ₋out[4, 2, k] = t72 - ˍ₋out[1, 3, k] = t86 - ˍ₋out[2, 3, k] = t97 - ˍ₋out[3, 3, k] = t105 - ˍ₋out[4, 3, k] = t110 - ˍ₋out[1, 4, k] = t124 - ˍ₋out[2, 4, k] = t135 - ˍ₋out[3, 4, k] = t143 - ˍ₋out[4, 4, k] = t148 - nothing - end - end - end - end -end - -function tensor_inverse4(A::AbstractArray{T, 3}) where {T} - out = similar(A) - - tensor_inverse4!(out, A) - - out -end - -function tensor_inverse4!(out::AbstractArray{T, 3}, A::AbstractArray{T, 3}) where {T} - @assert size(A, 1) == size(A, 2) == 4 - @assert size(A) == size(out) - - backend = networkbackend(out) - inv44! = inv44_kernel!(backend) - - inv44!(out, A, ndrange = size(A, 3)) - - nothing -end - -function ChainRulesCore.rrule(::typeof(tensor_inverse4), A::AT) where { - T, AT <: AbstractArray{T, 3}} - out = tensor_inverse4(A) - - function tensor_inverse_pullback(out_diff) - _tensor_inverse_pullback(out, out_diff) - end - out, tensor_inverse_pullback -end diff --git a/src/kernels/inverses/inverse_5x5.jl b/src/kernels/inverses/inverse_5x5.jl deleted file mode 100644 index 7686f7a49..000000000 --- a/src/kernels/inverses/inverse_5x5.jl +++ /dev/null @@ -1,484 +0,0 @@ -@kernel function inv55_kernel!(ˍ₋out, A) - k = @index(Global) - @inbounds begin - begin - begin - t1 = -1 - t2 = (*)((getindex)(A, 4, 2, k), (getindex)(A, 5, 5, k)) - t3 = (*)((*)(t1, (getindex)(A, 5, 2, k)), (getindex)(A, 4, 5, k)) - t4 = (+)(t2, t3) - t5 = (*)((*)(t1, (getindex)(A, 3, 3, k)), t4) - t6 = (*)((*)(t1, (getindex)(A, 5, 3, k)), (getindex)(A, 4, 5, k)) - t7 = (*)((getindex)(A, 5, 5, k), (getindex)(A, 4, 3, k)) - t8 = (+)(t6, t7) - t9 = (*)(t8, (getindex)(A, 3, 2, k)) - t10 = (*)((*)(t1, (getindex)(A, 5, 2, k)), (getindex)(A, 4, 3, k)) - t11 = (*)((getindex)(A, 4, 2, k), (getindex)(A, 5, 3, k)) - t12 = (+)(t10, t11) - t13 = (*)(t12, (getindex)(A, 3, 5, k)) - t14 = (+)((+)(t5, t9), t13) - t15 = (*)(t14, (getindex)(A, 2, 4, k)) - t16 = (*)((*)(t1, (getindex)(A, 4, 4, k)), (getindex)(A, 5, 3, k)) - t17 = (*)((getindex)(A, 5, 4, k), (getindex)(A, 4, 3, k)) - t18 = (+)(t16, t17) - t19 = (*)((getindex)(A, 3, 5, k), t18) - t20 = (*)((*)(t1, t8), (getindex)(A, 3, 4, k)) - t21 = (*)((getindex)(A, 4, 4, k), (getindex)(A, 5, 5, k)) - t22 = (*)((*)(t1, (getindex)(A, 5, 4, k)), (getindex)(A, 4, 5, k)) - t23 = (+)(t21, t22) - t24 = (*)((getindex)(A, 3, 3, k), t23) - t25 = (+)((+)(t19, t20), t24) - t26 = (*)(t25, (getindex)(A, 2, 2, k)) - t27 = (*)(t18, (getindex)(A, 3, 2, k)) - t28 = (*)(t12, (getindex)(A, 3, 4, k)) - t29 = (*)((getindex)(A, 4, 2, k), (getindex)(A, 5, 4, k)) - t30 = (*)((*)(t1, (getindex)(A, 4, 4, k)), (getindex)(A, 5, 2, k)) - t31 = (+)(t29, t30) - t32 = (*)((*)(t1, (getindex)(A, 3, 3, k)), t31) - t33 = (+)((+)(t27, t28), t32) - t34 = (*)((*)(t1, t33), (getindex)(A, 2, 5, k)) - t35 = (*)((*)(t1, (getindex)(A, 3, 4, k)), t4) - t36 = (*)((getindex)(A, 3, 5, k), t31) - t37 = (*)(t23, (getindex)(A, 3, 2, k)) - t38 = (+)((+)(t35, t36), t37) - t39 = (*)((*)(t1, (getindex)(A, 2, 3, k)), t38) - t40 = (+)((+)((+)(t15, t26), t34), t39) - t41 = (*)((getindex)(A, 5, 3, k), (getindex)(A, 4, 1, k)) - t42 = (*)((*)(t1, (getindex)(A, 5, 1, k)), (getindex)(A, 4, 3, k)) - t43 = (+)(t41, t42) - t44 = (*)((getindex)(A, 3, 5, k), t43) - t45 = (*)((*)(t1, (getindex)(A, 5, 1, k)), (getindex)(A, 4, 5, k)) - t46 = (*)((getindex)(A, 5, 5, k), (getindex)(A, 4, 1, k)) - t47 = (+)(t45, t46) - t48 = (*)((*)(t1, (getindex)(A, 3, 3, k)), t47) - t49 = (*)((getindex)(A, 3, 1, k), t8) - t50 = (+)((+)(t44, t48), t49) - t51 = (*)((getindex)(A, 2, 4, k), t50) - t52 = (*)(t23, (getindex)(A, 3, 1, k)) - t53 = (*)((*)(t1, t47), (getindex)(A, 3, 4, k)) - t54 = (*)((getindex)(A, 5, 4, k), (getindex)(A, 4, 1, k)) - t55 = (*)((*)(t1, (getindex)(A, 4, 4, k)), (getindex)(A, 5, 1, k)) - t56 = (+)(t54, t55) - t57 = (*)((getindex)(A, 3, 5, k), t56) - t58 = (+)((+)(t52, t53), t57) - t59 = (*)((*)(t1, (getindex)(A, 2, 3, k)), t58) - t60 = (*)((getindex)(A, 3, 4, k), t43) - t61 = (*)((*)(t1, (getindex)(A, 3, 3, k)), t56) - t62 = (*)((getindex)(A, 3, 1, k), t18) - t63 = (+)((+)(t60, t61), t62) - t64 = (*)((*)(t1, t63), (getindex)(A, 2, 5, k)) - t65 = (*)((getindex)(A, 2, 1, k), t25) - t66 = (+)((+)((+)(t51, t59), t64), t65) - t67 = (*)((*)(t1, (getindex)(A, 1, 2, k)), t66) - t68 = (*)((*)(t1, (getindex)(A, 4, 2, k)), (getindex)(A, 5, 1, k)) - t69 = (*)((getindex)(A, 5, 2, k), (getindex)(A, 4, 1, k)) - t70 = (+)(t68, t69) - t71 = (*)(t70, (getindex)(A, 3, 4, k)) - t72 = (*)((*)(t1, t56), (getindex)(A, 3, 2, k)) - t73 = (*)((getindex)(A, 3, 1, k), t31) - t74 = (+)((+)(t71, t72), t73) - t75 = (*)(t74, (getindex)(A, 2, 3, k)) - t76 = (*)(t33, (getindex)(A, 2, 1, k)) - t77 = (*)((*)(t1, t43), (getindex)(A, 3, 2, k)) - t78 = (*)(t12, (getindex)(A, 3, 1, k)) - t79 = (*)((getindex)(A, 3, 3, k), t70) - t80 = (+)((+)(t77, t78), t79) - t81 = (*)((*)(t1, (getindex)(A, 2, 4, k)), t80) - t82 = (*)((*)(t1, t63), (getindex)(A, 2, 2, k)) - t83 = (+)((+)((+)(t75, t76), t81), t82) - t84 = (*)(t83, (getindex)(A, 1, 5, k)) - t85 = (*)((getindex)(A, 1, 1, k), t40) - t86 = (*)((*)(t1, t74), (getindex)(A, 2, 5, k)) - t87 = (*)((getindex)(A, 2, 1, k), t38) - t88 = (*)((getindex)(A, 3, 1, k), t4) - t89 = (*)((getindex)(A, 3, 5, k), t70) - t90 = (*)((*)(t1, t47), (getindex)(A, 3, 2, k)) - t91 = (+)((+)(t88, t89), t90) - t92 = (*)(t91, (getindex)(A, 2, 4, k)) - t93 = (*)((*)(t1, t58), (getindex)(A, 2, 2, k)) - t94 = (+)((+)((+)(t86, t87), t92), t93) - t95 = (*)(t94, (getindex)(A, 1, 3, k)) - t96 = (*)(t91, (getindex)(A, 2, 3, k)) - t97 = (*)((*)(t1, (getindex)(A, 2, 5, k)), t80) - t98 = (*)((getindex)(A, 2, 1, k), t14) - t99 = (*)((*)(t1, (getindex)(A, 2, 2, k)), t50) - t100 = (+)((+)((+)(t96, t97), t98), t99) - t101 = (*)((*)(t1, (getindex)(A, 1, 4, k)), t100) - t102 = (+)((+)((+)((+)(t67, t84), t85), t95), t101) - t103 = (/)(t40, t102) - t104 = (*)((*)(t1, (getindex)(A, 2, 4, k)), t50) - t105 = (*)((getindex)(A, 2, 3, k), t58) - t106 = (*)(t63, (getindex)(A, 2, 5, k)) - t107 = (*)((*)(t1, (getindex)(A, 2, 1, k)), t25) - t108 = (+)((+)((+)(t104, t105), t106), t107) - t109 = (/)(t108, t102) - t110 = (/)(t94, t102) - t111 = (*)((*)(t1, t91), (getindex)(A, 2, 3, k)) - t112 = (*)((getindex)(A, 2, 5, k), t80) - t113 = (*)((*)(t1, (getindex)(A, 2, 1, k)), t14) - t114 = (*)((getindex)(A, 2, 2, k), t50) - t115 = (+)((+)((+)(t111, t112), t113), t114) - t116 = (/)(t115, t102) - t117 = (/)(t83, t102) - t118 = (*)((*)(t1, (getindex)(A, 1, 2, k)), t25) - t119 = (*)((*)(t1, (getindex)(A, 1, 4, k)), t14) - t120 = (*)((getindex)(A, 1, 3, k), t38) - t121 = (*)(t33, (getindex)(A, 1, 5, k)) - t122 = (+)((+)((+)(t118, t119), t120), t121) - t123 = (/)(t122, t102) - t124 = (*)((*)(t1, t63), (getindex)(A, 1, 5, k)) - t125 = (*)((*)(t1, (getindex)(A, 1, 3, k)), t58) - t126 = (*)((getindex)(A, 1, 4, k), t50) - t127 = (*)((getindex)(A, 1, 1, k), t25) - t128 = (+)((+)((+)(t124, t125), t126), t127) - t129 = (/)(t128, t102) - t130 = (*)(t74, (getindex)(A, 1, 5, k)) - t131 = (*)((*)(t1, (getindex)(A, 1, 1, k)), t38) - t132 = (*)((*)(t1, (getindex)(A, 1, 4, k)), t91) - t133 = (*)((getindex)(A, 1, 2, k), t58) - t134 = (+)((+)((+)(t130, t131), t132), t133) - t135 = (/)(t134, t102) - t136 = (*)((*)(t1, (getindex)(A, 1, 2, k)), t50) - t137 = (*)(t91, (getindex)(A, 1, 3, k)) - t138 = (*)((*)(t1, (getindex)(A, 1, 5, k)), t80) - t139 = (*)((getindex)(A, 1, 1, k), t14) - t140 = (+)((+)((+)(t136, t137), t138), t139) - t141 = (/)(t140, t102) - t142 = (*)((*)(t1, (getindex)(A, 1, 1, k)), t33) - t143 = (*)((getindex)(A, 1, 2, k), t63) - t144 = (*)((*)(t1, t74), (getindex)(A, 1, 3, k)) - t145 = (*)((getindex)(A, 1, 4, k), t80) - t146 = (+)((+)((+)(t142, t143), t144), t145) - t147 = (/)(t146, t102) - t148 = (*)(t12, (getindex)(A, 2, 5, k)) - t149 = (*)((*)(t1, (getindex)(A, 2, 3, k)), t4) - t150 = (*)(t8, (getindex)(A, 2, 2, k)) - t151 = (+)((+)(t148, t149), t150) - t152 = (*)((getindex)(A, 1, 4, k), t151) - t153 = (*)(t18, (getindex)(A, 2, 5, k)) - t154 = (*)(t23, (getindex)(A, 2, 3, k)) - t155 = (*)((*)(t1, t8), (getindex)(A, 2, 4, k)) - t156 = (+)((+)(t153, t154), t155) - t157 = (*)((getindex)(A, 1, 2, k), t156) - t158 = (*)((getindex)(A, 2, 5, k), t31) - t159 = (*)(t23, (getindex)(A, 2, 2, k)) - t160 = (*)((*)(t1, t4), (getindex)(A, 2, 4, k)) - t161 = (+)((+)(t158, t159), t160) - t162 = (*)((*)(t1, (getindex)(A, 1, 3, k)), t161) - t163 = (*)(t12, (getindex)(A, 2, 4, k)) - t164 = (*)(t18, (getindex)(A, 2, 2, k)) - t165 = (*)((*)(t1, (getindex)(A, 2, 3, k)), t31) - t166 = (+)((+)(t163, t164), t165) - t167 = (*)((*)(t1, t166), (getindex)(A, 1, 5, k)) - t168 = (+)((+)((+)(t152, t157), t162), t167) - t169 = (/)(t168, t102) - t170 = (*)((getindex)(A, 2, 1, k), t8) - t171 = (*)((*)(t1, (getindex)(A, 2, 3, k)), t47) - t172 = (*)((getindex)(A, 2, 5, k), t43) - t173 = (+)((+)(t170, t171), t172) - t174 = (*)((*)(t1, (getindex)(A, 1, 4, k)), t173) - t175 = (*)((*)(t1, t47), (getindex)(A, 2, 4, k)) - t176 = (*)(t56, (getindex)(A, 2, 5, k)) - t177 = (*)(t23, (getindex)(A, 2, 1, k)) - t178 = (+)((+)(t175, t176), t177) - t179 = (*)((getindex)(A, 1, 3, k), t178) - t180 = (*)((*)(t1, (getindex)(A, 1, 1, k)), t156) - t181 = (*)(t43, (getindex)(A, 2, 4, k)) - t182 = (*)((getindex)(A, 2, 1, k), t18) - t183 = (*)((*)(t1, t56), (getindex)(A, 2, 3, k)) - t184 = (+)((+)(t181, t182), t183) - t185 = (*)(t184, (getindex)(A, 1, 5, k)) - t186 = (+)((+)((+)(t174, t179), t180), t185) - t187 = (/)(t186, t102) - t188 = (*)((getindex)(A, 1, 1, k), t161) - t189 = (*)(t70, (getindex)(A, 2, 5, k)) - t190 = (*)((getindex)(A, 2, 1, k), t4) - t191 = (*)((*)(t1, t47), (getindex)(A, 2, 2, k)) - t192 = (+)((+)(t189, t190), t191) - t193 = (*)((getindex)(A, 1, 4, k), t192) - t194 = (*)((getindex)(A, 2, 1, k), t31) - t195 = (*)((*)(t1, t56), (getindex)(A, 2, 2, k)) - t196 = (*)(t70, (getindex)(A, 2, 4, k)) - t197 = (+)((+)(t194, t195), t196) - t198 = (*)((*)(t1, (getindex)(A, 1, 5, k)), t197) - t199 = (*)((*)(t1, (getindex)(A, 1, 2, k)), t178) - t200 = (+)((+)((+)(t188, t193), t198), t199) - t201 = (/)(t200, t102) - t202 = (*)((*)(t1, (getindex)(A, 1, 1, k)), t151) - t203 = (*)((getindex)(A, 2, 3, k), t70) - t204 = (*)(t12, (getindex)(A, 2, 1, k)) - t205 = (*)((*)(t1, t43), (getindex)(A, 2, 2, k)) - t206 = (+)((+)(t203, t204), t205) - t207 = (*)((getindex)(A, 1, 5, k), t206) - t208 = (*)((*)(t1, (getindex)(A, 1, 3, k)), t192) - t209 = (*)((getindex)(A, 1, 2, k), t173) - t210 = (+)((+)((+)(t202, t207), t208), t209) - t211 = (/)(t210, t102) - t212 = (*)((getindex)(A, 1, 1, k), t166) - t213 = (*)((*)(t1, (getindex)(A, 1, 2, k)), t184) - t214 = (*)((*)(t1, (getindex)(A, 1, 4, k)), t206) - t215 = (*)((getindex)(A, 1, 3, k), t197) - t216 = (+)((+)((+)(t212, t213), t214), t215) - t217 = (/)(t216, t102) - t218 = (*)((getindex)(A, 5, 5, k), (getindex)(A, 3, 4, k)) - t219 = (*)((*)(t1, (getindex)(A, 3, 5, k)), (getindex)(A, 5, 4, k)) - t220 = (+)(t218, t219) - t221 = (*)((getindex)(A, 2, 3, k), t220) - t222 = (*)((getindex)(A, 3, 3, k), (getindex)(A, 5, 5, k)) - t223 = (*)((*)(t1, (getindex)(A, 3, 5, k)), (getindex)(A, 5, 3, k)) - t224 = (+)(t222, t223) - t225 = (*)((*)(t1, t224), (getindex)(A, 2, 4, k)) - t226 = (*)((*)(t1, (getindex)(A, 5, 3, k)), (getindex)(A, 3, 4, k)) - t227 = (*)((getindex)(A, 3, 3, k), (getindex)(A, 5, 4, k)) - t228 = (+)(t226, t227) - t229 = (*)(t228, (getindex)(A, 2, 5, k)) - t230 = (+)((+)(t221, t225), t229) - t231 = (*)((*)(t1, (getindex)(A, 1, 2, k)), t230) - t232 = (*)((*)(t1, (getindex)(A, 3, 5, k)), (getindex)(A, 5, 2, k)) - t233 = (*)((getindex)(A, 5, 5, k), (getindex)(A, 3, 2, k)) - t234 = (+)(t232, t233) - t235 = (*)((*)(t1, t234), (getindex)(A, 2, 4, k)) - t236 = (*)(t220, (getindex)(A, 2, 2, k)) - t237 = (*)((getindex)(A, 5, 4, k), (getindex)(A, 3, 2, k)) - t238 = (*)((*)(t1, (getindex)(A, 5, 2, k)), (getindex)(A, 3, 4, k)) - t239 = (+)(t237, t238) - t240 = (*)((getindex)(A, 2, 5, k), t239) - t241 = (+)((+)(t235, t236), t240) - t242 = (*)((getindex)(A, 1, 3, k), t241) - t243 = (*)((getindex)(A, 5, 3, k), (getindex)(A, 3, 2, k)) - t244 = (*)((*)(t1, (getindex)(A, 3, 3, k)), (getindex)(A, 5, 2, k)) - t245 = (+)(t243, t244) - t246 = (*)(t245, (getindex)(A, 2, 5, k)) - t247 = (*)(t224, (getindex)(A, 2, 2, k)) - t248 = (*)((*)(t1, (getindex)(A, 2, 3, k)), t234) - t249 = (+)((+)(t246, t247), t248) - t250 = (*)((*)(t1, (getindex)(A, 1, 4, k)), t249) - t251 = (*)(t228, (getindex)(A, 2, 2, k)) - t252 = (*)((*)(t1, (getindex)(A, 2, 3, k)), t239) - t253 = (*)(t245, (getindex)(A, 2, 4, k)) - t254 = (+)((+)(t251, t252), t253) - t255 = (*)(t254, (getindex)(A, 1, 5, k)) - t256 = (+)((+)((+)(t231, t242), t250), t255) - t257 = (/)(t256, t102) - t258 = (*)((getindex)(A, 5, 4, k), (getindex)(A, 3, 1, k)) - t259 = (*)((*)(t1, (getindex)(A, 5, 1, k)), (getindex)(A, 3, 4, k)) - t260 = (+)(t258, t259) - t261 = (*)((*)(t1, t260), (getindex)(A, 2, 3, k)) - t262 = (*)(t228, (getindex)(A, 2, 1, k)) - t263 = (*)((getindex)(A, 3, 1, k), (getindex)(A, 5, 3, k)) - t264 = (*)((*)(t1, (getindex)(A, 3, 3, k)), (getindex)(A, 5, 1, k)) - t265 = (+)(t263, t264) - t266 = (*)(t265, (getindex)(A, 2, 4, k)) - t267 = (+)((+)(t261, t262), t266) - t268 = (*)((*)(t1, t267), (getindex)(A, 1, 5, k)) - t269 = (*)(t230, (getindex)(A, 1, 1, k)) - t270 = (*)((*)(t1, (getindex)(A, 3, 5, k)), (getindex)(A, 5, 1, k)) - t271 = (*)((getindex)(A, 3, 1, k), (getindex)(A, 5, 5, k)) - t272 = (+)(t270, t271) - t273 = (*)((*)(t1, (getindex)(A, 2, 3, k)), t272) - t274 = (*)((getindex)(A, 2, 1, k), t224) - t275 = (*)(t265, (getindex)(A, 2, 5, k)) - t276 = (+)((+)(t273, t274), t275) - t277 = (*)((getindex)(A, 1, 4, k), t276) - t278 = (*)((*)(t1, t272), (getindex)(A, 2, 4, k)) - t279 = (*)((getindex)(A, 2, 1, k), t220) - t280 = (*)(t260, (getindex)(A, 2, 5, k)) - t281 = (+)((+)(t278, t279), t280) - t282 = (*)((*)(t1, t281), (getindex)(A, 1, 3, k)) - t283 = (+)((+)((+)(t268, t269), t277), t282) - t284 = (/)(t283, t102) - t285 = (*)((*)(t1, (getindex)(A, 1, 1, k)), t241) - t286 = (*)((getindex)(A, 1, 2, k), t281) - t287 = (*)((*)(t1, t260), (getindex)(A, 2, 2, k)) - t288 = (*)((getindex)(A, 3, 1, k), (getindex)(A, 5, 2, k)) - t289 = (*)((*)(t1, (getindex)(A, 5, 1, k)), (getindex)(A, 3, 2, k)) - t290 = (+)(t288, t289) - t291 = (*)(t290, (getindex)(A, 2, 4, k)) - t292 = (*)((getindex)(A, 2, 1, k), t239) - t293 = (+)((+)(t287, t291), t292) - t294 = (*)(t293, (getindex)(A, 1, 5, k)) - t295 = (*)((*)(t1, t272), (getindex)(A, 2, 2, k)) - t296 = (*)(t290, (getindex)(A, 2, 5, k)) - t297 = (*)((getindex)(A, 2, 1, k), t234) - t298 = (+)((+)(t295, t296), t297) - t299 = (*)((*)(t1, (getindex)(A, 1, 4, k)), t298) - t300 = (+)((+)((+)(t285, t286), t294), t299) - t301 = (/)(t300, t102) - t302 = (*)((getindex)(A, 1, 3, k), t298) - t303 = (*)((*)(t1, (getindex)(A, 1, 2, k)), t276) - t304 = (*)((getindex)(A, 1, 1, k), t249) - t305 = (*)(t245, (getindex)(A, 2, 1, k)) - t306 = (*)((*)(t1, t265), (getindex)(A, 2, 2, k)) - t307 = (*)((getindex)(A, 2, 3, k), t290) - t308 = (+)((+)(t305, t306), t307) - t309 = (*)((*)(t1, t308), (getindex)(A, 1, 5, k)) - t310 = (+)((+)((+)(t302, t303), t304), t309) - t311 = (/)(t310, t102) - t312 = (*)((getindex)(A, 1, 4, k), t308) - t313 = (*)((getindex)(A, 1, 2, k), t267) - t314 = (*)((*)(t1, t293), (getindex)(A, 1, 3, k)) - t315 = (*)((*)(t1, (getindex)(A, 1, 1, k)), t254) - t316 = (+)((+)((+)(t312, t313), t314), t315) - t317 = (/)(t316, t102) - t318 = (*)((getindex)(A, 4, 5, k), (getindex)(A, 3, 2, k)) - t319 = (*)((*)(t1, (getindex)(A, 4, 2, k)), (getindex)(A, 3, 5, k)) - t320 = (+)(t318, t319) - t321 = (*)((*)(t1, t320), (getindex)(A, 2, 3, k)) - t322 = (*)((getindex)(A, 4, 3, k), (getindex)(A, 3, 2, k)) - t323 = (*)((*)(t1, (getindex)(A, 3, 3, k)), (getindex)(A, 4, 2, k)) - t324 = (+)(t322, t323) - t325 = (*)(t324, (getindex)(A, 2, 5, k)) - t326 = (*)((*)(t1, (getindex)(A, 3, 5, k)), (getindex)(A, 4, 3, k)) - t327 = (*)((getindex)(A, 3, 3, k), (getindex)(A, 4, 5, k)) - t328 = (+)(t326, t327) - t329 = (*)(t328, (getindex)(A, 2, 2, k)) - t330 = (+)((+)(t321, t325), t329) - t331 = (*)((getindex)(A, 1, 4, k), t330) - t332 = (*)((getindex)(A, 3, 3, k), (getindex)(A, 4, 4, k)) - t333 = (*)((*)(t1, (getindex)(A, 4, 3, k)), (getindex)(A, 3, 4, k)) - t334 = (+)(t332, t333) - t335 = (*)(t334, (getindex)(A, 2, 2, k)) - t336 = (*)((getindex)(A, 4, 4, k), (getindex)(A, 3, 2, k)) - t337 = (*)((*)(t1, (getindex)(A, 4, 2, k)), (getindex)(A, 3, 4, k)) - t338 = (+)(t336, t337) - t339 = (*)((*)(t1, (getindex)(A, 2, 3, k)), t338) - t340 = (*)(t324, (getindex)(A, 2, 4, k)) - t341 = (+)((+)(t335, t339), t340) - t342 = (*)((*)(t1, t341), (getindex)(A, 1, 5, k)) - t343 = (*)((getindex)(A, 3, 4, k), (getindex)(A, 4, 5, k)) - t344 = (*)((*)(t1, (getindex)(A, 3, 5, k)), (getindex)(A, 4, 4, k)) - t345 = (+)(t343, t344) - t346 = (*)((getindex)(A, 2, 3, k), t345) - t347 = (*)((*)(t1, t328), (getindex)(A, 2, 4, k)) - t348 = (*)(t334, (getindex)(A, 2, 5, k)) - t349 = (+)((+)(t346, t347), t348) - t350 = (*)((getindex)(A, 1, 2, k), t349) - t351 = (*)((*)(t1, t320), (getindex)(A, 2, 4, k)) - t352 = (*)((getindex)(A, 2, 5, k), t338) - t353 = (*)(t345, (getindex)(A, 2, 2, k)) - t354 = (+)((+)(t351, t352), t353) - t355 = (*)((*)(t1, (getindex)(A, 1, 3, k)), t354) - t356 = (+)((+)((+)(t331, t342), t350), t355) - t357 = (/)(t356, t102) - t358 = (*)((getindex)(A, 3, 1, k), (getindex)(A, 4, 5, k)) - t359 = (*)((*)(t1, (getindex)(A, 3, 5, k)), (getindex)(A, 4, 1, k)) - t360 = (+)(t358, t359) - t361 = (*)((*)(t1, t360), (getindex)(A, 2, 3, k)) - t362 = (*)((getindex)(A, 2, 1, k), t328) - t363 = (*)((getindex)(A, 3, 1, k), (getindex)(A, 4, 3, k)) - t364 = (*)((*)(t1, (getindex)(A, 3, 3, k)), (getindex)(A, 4, 1, k)) - t365 = (+)(t363, t364) - t366 = (*)(t365, (getindex)(A, 2, 5, k)) - t367 = (+)((+)(t361, t362), t366) - t368 = (*)((*)(t1, (getindex)(A, 1, 4, k)), t367) - t369 = (*)((getindex)(A, 2, 1, k), t345) - t370 = (*)((*)(t1, t360), (getindex)(A, 2, 4, k)) - t371 = (*)((getindex)(A, 4, 4, k), (getindex)(A, 3, 1, k)) - t372 = (*)((*)(t1, (getindex)(A, 4, 1, k)), (getindex)(A, 3, 4, k)) - t373 = (+)(t371, t372) - t374 = (*)(t373, (getindex)(A, 2, 5, k)) - t375 = (+)((+)(t369, t370), t374) - t376 = (*)((getindex)(A, 1, 3, k), t375) - t377 = (*)((*)(t1, (getindex)(A, 2, 3, k)), t373) - t378 = (*)(t365, (getindex)(A, 2, 4, k)) - t379 = (*)((getindex)(A, 2, 1, k), t334) - t380 = (+)((+)(t377, t378), t379) - t381 = (*)(t380, (getindex)(A, 1, 5, k)) - t382 = (*)((*)(t1, (getindex)(A, 1, 1, k)), t349) - t383 = (+)((+)((+)(t368, t376), t381), t382) - t384 = (/)(t383, t102) - t385 = (*)((getindex)(A, 2, 1, k), t338) - t386 = (*)((*)(t1, (getindex)(A, 4, 1, k)), (getindex)(A, 3, 2, k)) - t387 = (*)((getindex)(A, 4, 2, k), (getindex)(A, 3, 1, k)) - t388 = (+)(t386, t387) - t389 = (*)(t388, (getindex)(A, 2, 4, k)) - t390 = (*)((*)(t1, t373), (getindex)(A, 2, 2, k)) - t391 = (+)((+)(t385, t389), t390) - t392 = (*)((*)(t1, t391), (getindex)(A, 1, 5, k)) - t393 = (*)(t320, (getindex)(A, 2, 1, k)) - t394 = (*)((*)(t1, t360), (getindex)(A, 2, 2, k)) - t395 = (*)((getindex)(A, 2, 5, k), t388) - t396 = (+)((+)(t393, t394), t395) - t397 = (*)((getindex)(A, 1, 4, k), t396) - t398 = (*)((getindex)(A, 1, 1, k), t354) - t399 = (*)((*)(t1, (getindex)(A, 1, 2, k)), t375) - t400 = (+)((+)((+)(t392, t397), t398), t399) - t401 = (/)(t400, t102) - t402 = (*)((getindex)(A, 1, 2, k), t367) - t403 = (*)((*)(t1, t330), (getindex)(A, 1, 1, k)) - t404 = (*)((getindex)(A, 2, 3, k), t388) - t405 = (*)((*)(t1, t365), (getindex)(A, 2, 2, k)) - t406 = (*)(t324, (getindex)(A, 2, 1, k)) - t407 = (+)((+)(t404, t405), t406) - t408 = (*)((getindex)(A, 1, 5, k), t407) - t409 = (*)((*)(t1, (getindex)(A, 1, 3, k)), t396) - t410 = (+)((+)((+)(t402, t403), t408), t409) - t411 = (/)(t410, t102) - t412 = (*)((getindex)(A, 1, 3, k), t391) - t413 = (*)((*)(t1, (getindex)(A, 1, 2, k)), t380) - t414 = (*)((*)(t1, (getindex)(A, 1, 4, k)), t407) - t415 = (*)((getindex)(A, 1, 1, k), t341) - t416 = (+)((+)((+)(t412, t413), t414), t415) - t417 = (/)(t416, t102) - @inbounds begin - ˍ₋out[1, 1, k] = t103 - ˍ₋out[2, 1, k] = t109 - ˍ₋out[3, 1, k] = t110 - ˍ₋out[4, 1, k] = t116 - ˍ₋out[5, 1, k] = t117 - ˍ₋out[1, 2, k] = t123 - ˍ₋out[2, 2, k] = t129 - ˍ₋out[3, 2, k] = t135 - ˍ₋out[4, 2, k] = t141 - ˍ₋out[5, 2, k] = t147 - ˍ₋out[1, 3, k] = t169 - ˍ₋out[2, 3, k] = t187 - ˍ₋out[3, 3, k] = t201 - ˍ₋out[4, 3, k] = t211 - ˍ₋out[5, 3, k] = t217 - ˍ₋out[1, 4, k] = t257 - ˍ₋out[2, 4, k] = t284 - ˍ₋out[3, 4, k] = t301 - ˍ₋out[4, 4, k] = t311 - ˍ₋out[5, 4, k] = t317 - ˍ₋out[1, 5, k] = t357 - ˍ₋out[2, 5, k] = t384 - ˍ₋out[3, 5, k] = t401 - ˍ₋out[4, 5, k] = t411 - ˍ₋out[5, 5, k] = t417 - nothing - end - end - end - end -end - -function tensor_inverse5(A::AbstractArray{T, 3}) where {T} - out = similar(A) - - tensor_inverse5!(out, A) - - out -end - -function tensor_inverse5!(out::AbstractArray{T, 3}, A::AbstractArray{T, 3}) where {T} - @assert size(A, 1) == size(A, 2) == 5 - @assert size(A) == size(out) - - backend = networkbackend(out) - inv55! = inv55_kernel!(backend) - - inv55!(out, A, ndrange = size(A, 3)) - - nothing -end - -function ChainRulesCore.rrule(::typeof(tensor_inverse5), A::AT) where { - T, AT <: AbstractArray{T, 3}} - out = tensor_inverse5(A) - - function tensor_inverse_pullback(out_diff) - _tensor_inverse_pullback(out, out_diff) - end - out, tensor_inverse_pullback -end diff --git a/src/kernels/inverses/tensor_cayley.jl b/src/kernels/inverses/tensor_cayley.jl deleted file mode 100644 index 473641534..000000000 --- a/src/kernels/inverses/tensor_cayley.jl +++ /dev/null @@ -1,19 +0,0 @@ -function tensor_cayley2(A::AbstractArray{T, 3}) where {T} - one_tensor = init_output(A) - tensor_tensor_mul(one_tensor - A, tensor_inverse2(one_tensor + A)) -end - -function tensor_cayley3(A::AbstractArray{T, 3}) where {T} - one_tensor = init_output(A) - tensor_tensor_mul(one_tensor - A, tensor_inverse3(one_tensor + A)) -end - -function tensor_cayley4(A::AbstractArray{T, 3}) where {T} - one_tensor = init_output(A) - tensor_tensor_mul(one_tensor - A, tensor_inverse4(one_tensor + A)) -end - -function tensor_cayley5(A::AbstractArray{T, 3}) where {T} - one_tensor = init_output(A) - tensor_tensor_mul(one_tensor - A, tensor_inverse5(one_tensor + A)) -end diff --git a/src/kernels/inverses/tensor_inverse_pullback.jl b/src/kernels/inverses/tensor_inverse_pullback.jl deleted file mode 100644 index 1a0ebf1bd..000000000 --- a/src/kernels/inverses/tensor_inverse_pullback.jl +++ /dev/null @@ -1,10 +0,0 @@ -# The pullback of `tensor_inverse2` … `tensor_inverse5`, whose `rrule`s -# `scripts/utilities/inverse_generator.jl` writes: for the inverse `out`, slice by slice, -# `Ā = -out' * out_diff * out'`. -function _tensor_inverse_pullback(out::AbstractArray{T, 3}, out_diff) where {T} - out_diff = _seed(T, out_diff) - out_diff isa AbstractZero && return NoTangent(), ZeroTangent() - - NoTangent(), - -tensor_transpose_tensor_mul(out, tensor_tensor_transpose_mul(out_diff, out)) -end diff --git a/src/layers/volume_preserving_attention.jl b/src/layers/volume_preserving_attention.jl index 18a6a4e0c..5d299591c 100644 --- a/src/layers/volume_preserving_attention.jl +++ b/src/layers/volume_preserving_attention.jl @@ -19,9 +19,7 @@ Applying a layer of type `VolumePreservingAttention` does the following: # Implementation -The layer's `seq_length` chooses the activation: a `seq_length` of 2, 3, 4 or 5 uses the closed form -`tensor_cayley2` to `tensor_cayley5`, whose inverse has been computed symbolically. Every other -`seq_length`, `0` included, uses [`tensor_cayley`](@ref), which solves ``(\mathbb{I} + C)\Lambda = +The Cayley transform is computed by [`tensor_cayley`](@ref), which solves ``(\mathbb{I} + C)\Lambda = \mathbb{I} - C`` by row-pivoted Gauss–Jordan elimination at any sequence length, on the CPU and on a device, and in the element type of the data. """ @@ -40,30 +38,6 @@ function orthonormal_activation_cayley(::VolumePreservingAttention, A::AbstractA tensor_cayley(A) end -function orthonormal_activation_cayley( - ::VolumePreservingAttention{M, M, ScalarProductType, 2}, - A::AbstractArray{T, 3}) where {T, M, ScalarProductType} - tensor_cayley2(A) -end - -function orthonormal_activation_cayley( - ::VolumePreservingAttention{M, M, ScalarProductType, 3}, - A::AbstractArray{T, 3}) where {T, M, ScalarProductType} - tensor_cayley3(A) -end - -function orthonormal_activation_cayley( - ::VolumePreservingAttention{M, M, ScalarProductType, 4}, - A::AbstractArray{T, 3}) where {T, M, ScalarProductType} - tensor_cayley4(A) -end - -function orthonormal_activation_cayley( - ::VolumePreservingAttention{M, M, ScalarProductType, 5}, - A::AbstractArray{T, 3}) where {T, M, ScalarProductType} - tensor_cayley5(A) -end - function parameterlength(::VolumePreservingAttention{M, M, :skew_sym}) where {M} M * (M-1) ÷ 2 end diff --git a/test/kernels/inverses/test_cayley_transforms.jl b/test/kernels/inverses/test_cayley_transforms.jl deleted file mode 100644 index bb80dd15a..000000000 --- a/test/kernels/inverses/test_cayley_transforms.jl +++ /dev/null @@ -1,28 +0,0 @@ -using GeometricMachineLearning: tensor_cayley4, tensor_cayley3 -using Test -import Random - -Random.seed!(1234) - -function test_orthonormal(A::AbstractMatrix) - @test A' * A ≈ one(A) -end - -function test_orthonormal(A::AbstractArray{T, 3}) where {T} - for i in axes(A, 3) - A_temp = @view A[:, :, i] - test_orthonormal(A_temp) - end -end - -function test_cayley_transforms(third_dim::Int = 10) - A₄ = rand(4, 4, third_dim) - A₃ = rand(3, 3, third_dim) - # asymmetrize - A₄ = A₄ - permutedims(A₄, (2, 1, 3)) - A₃ = A₃ - permutedims(A₃, (2, 1, 3)) - test_orthonormal(tensor_cayley4(A₄)) - test_orthonormal(tensor_cayley3(A₃)) -end - -test_cayley_transforms() diff --git a/test/kernels/kernel_pullbacks.jl b/test/kernels/kernel_pullbacks.jl index 17c75c9d0..4a0b117f0 100644 --- a/test/kernels/kernel_pullbacks.jl +++ b/test/kernels/kernel_pullbacks.jl @@ -1,12 +1,8 @@ using GeometricMachineLearning: tensor_tensor_mul, tensor_transpose_tensor_mul, - tensor_mat_skew_sym_assign, tensor_inverse2, - tensor_inverse3, tensor_inverse4, tensor_inverse5, - init_output, tensor_cayley, - _mul + tensor_mat_skew_sym_assign, tensor_cayley, _mul using ChainRulesCore: rrule, AbstractZero, NoTangent, ZeroTangent, @thunk using ChainRulesTestUtils using FiniteDifferences: central_fdm, j′vp -using LinearAlgebra: I using Test import Random @@ -14,9 +10,6 @@ Random.seed!(1234) const verbose = false -# `n × n` slices that are far from singular, for the inverse rules. -invertible(T, n, k) = cat((rand(T, n, n) + n * I for _ in 1:k)...; dims = 3) - # `skew_batch`, a skew batch for the Cayley rule, whose slices are square. include("../helpers/cayley.jl") @@ -30,7 +23,6 @@ function main(first_dim, second_dim, third_dim, third_tensor_dim) rand(first_dim, first_dim)) test_rrule(_mul, rand(first_dim, second_dim)', rand(first_dim, third_dim, third_tensor_dim)) test_rrule(_mul, rand(first_dim, second_dim)', rand(first_dim, third_dim)) - test_rrule(init_output, rand(first_dim, first_dim, third_tensor_dim)) test_rrule(tensor_cayley, skew_batch(Float64, first_dim, third_tensor_dim)) end @@ -47,22 +39,12 @@ for _ in 1:num_tests verbose && println() end -for (n, f) in enumerate(( - tensor_inverse2, tensor_inverse3, tensor_inverse4, tensor_inverse5)) - test_rrule(f, invertible(Float64, n + 1, 3)) -end - # Every `rrule` of `src/kernels/` with a primal of element type `T`, and its arguments. function rules(T) ((tensor_tensor_mul, (rand(T, 3, 4, 2), rand(T, 4, 5, 2))), (tensor_transpose_tensor_mul, (rand(T, 4, 3, 2), rand(T, 4, 5, 2))), (tensor_mat_skew_sym_assign, (rand(T, 3, 4, 2), rand(T, 3, 3))), (_mul, (rand(T, 4, 3)', rand(T, 4, 5, 2))), - (tensor_inverse2, (invertible(T, 2, 2),)), - (tensor_inverse3, (invertible(T, 3, 2),)), - (tensor_inverse4, (invertible(T, 4, 2),)), - (tensor_inverse5, (invertible(T, 5, 2),)), - (init_output, (rand(T, 3, 3, 2),)), (tensor_cayley, (skew_batch(T, 3, 2),))) end @@ -84,7 +66,7 @@ end c̄ = pullback(S.(y)) @test c̄[1] isa NoTangent for c in Base.tail(c̄) - @test f === init_output ? c isa AbstractZero : eltype(c) === T + @test eltype(c) === T end end end diff --git a/test/kernels/tensor_cayley.jl b/test/kernels/tensor_cayley.jl deleted file mode 100644 index e30a55dab..000000000 --- a/test/kernels/tensor_cayley.jl +++ /dev/null @@ -1,47 +0,0 @@ -using GeometricMachineLearning: tensor_cayley2, tensor_cayley3, tensor_cayley4, - tensor_cayley5 -using Test -import Random - -Random.seed!(123) - -function check_orthonormal_property(B::AbstractArray) - for i in axes(B, 3) - B_temp = @view B[:, :, i] - @test B_temp * B_temp' ≈ one(B_temp) - end -end - -function test_tensor_cayley2(T::Type = Float32, third_dim::Int = 10) - A = rand(T, 2, 2, third_dim) - B = tensor_cayley2(A - permutedims(A, (2, 1, 3))) - check_orthonormal_property(B) -end - -function test_tensor_cayley3(T::Type = Float32, third_dim::Int = 10) - A = rand(T, 3, 3, third_dim) - B = tensor_cayley3(A - permutedims(A, (2, 1, 3))) - check_orthonormal_property(B) -end - -function test_tensor_cayley4(T::Type = Float32, third_dim::Int = 10) - A = rand(T, 4, 4, third_dim) - B = tensor_cayley4(A - permutedims(A, (2, 1, 3))) - check_orthonormal_property(B) -end - -function test_tensor_cayley5(T::Type = Float32, third_dim::Int = 10) - A = rand(T, 5, 5, third_dim) - B = tensor_cayley5(A - permutedims(A, (2, 1, 3))) - check_orthonormal_property(B) -end - -function check_all(T::Type, third_dim::Int = 10) - test_tensor_cayley2(T, third_dim) - test_tensor_cayley3(T, third_dim) - test_tensor_cayley4(T, third_dim) - test_tensor_cayley5(T, third_dim) -end - -check_all(Float32) -check_all(Float64) diff --git a/test/kernels/tensor_inverse.jl b/test/kernels/tensor_inverse.jl deleted file mode 100644 index 82a82ec34..000000000 --- a/test/kernels/tensor_inverse.jl +++ /dev/null @@ -1,123 +0,0 @@ -using GeometricMachineLearning: tensor_inverse2, tensor_inverse3, tensor_inverse4, - tensor_inverse5 -using Test -import Zygote -import Random - -Random.seed!(1234) - -function test55_inverse(k::Int = 10) - A = rand(5, 5, k) - - A_inv = tensor_inverse5(A) - for i in 1:k - @test inv(A[:, :, i]) ≈ A_inv[:, :, i] - end -end - -test55_inverse() - -function test55_inverse_pullback(k::Int = 10) - A = rand(5, 5, k) - - pullback_total = Zygote.pullback(tensor_inverse5, A) - - out_diff = rand(5, 5, k) - in_diff = pullback_total[2](out_diff)[1] - - for i in 1:k - pullback_i = Zygote.pullback(inv, A[:, :, i]) - - @test pullback_total[1][:, :, i] ≈ pullback_i[1] - @test in_diff[:, :, i] ≈ pullback_i[2](out_diff[:, :, i])[1] - end -end - -test55_inverse_pullback() - -function test44_inverse(k::Int = 10) - A = rand(4, 4, k) - - A_inv = tensor_inverse4(A) - for i in 1:k - @test inv(A[:, :, i]) ≈ A_inv[:, :, i] - end -end - -test44_inverse() - -function test44_inverse_pullback(k::Int = 10) - A = rand(4, 4, k) - - pullback_total = Zygote.pullback(tensor_inverse4, A) - - out_diff = rand(4, 4, k) - in_diff = pullback_total[2](out_diff)[1] - - for i in 1:k - pullback_i = Zygote.pullback(inv, A[:, :, i]) - - @test pullback_total[1][:, :, i] ≈ pullback_i[1] - @test in_diff[:, :, i] ≈ pullback_i[2](out_diff[:, :, i])[1] - end -end - -test44_inverse_pullback() - -function test33_inverse(k::Int = 10) - A = rand(3, 3, k) - - A_inv = tensor_inverse3(A) - for i in 1:k - @test inv(A[:, :, i]) ≈ A_inv[:, :, i] - end -end - -test33_inverse() - -function test33_inverse_pullback(k::Int = 10) - A = rand(3, 3, k) - - pullback_total = Zygote.pullback(tensor_inverse3, A) - - out_diff = rand(3, 3, k) - in_diff = pullback_total[2](out_diff)[1] - - for i in 1:k - pullback_i = Zygote.pullback(inv, A[:, :, i]) - - @test pullback_total[1][:, :, i] ≈ pullback_i[1] - @test in_diff[:, :, i] ≈ pullback_i[2](out_diff[:, :, i])[1] - end -end - -test33_inverse_pullback() - -function test22_inverse(k::Int = 10) - A = rand(2, 2, k) - - A_inv = tensor_inverse2(A) - for i in 1:k - @test inv(A[:, :, i]) ≈ A_inv[:, :, i] - end -end - -test22_inverse() - -function test22_inverse_pullback(k::Int = 10) - A = rand(2, 2, k) - - pullback_total = Zygote.pullback(tensor_inverse2, A) - - out_diff = rand(2, 2, k) - in_diff = pullback_total[2](out_diff)[1] - - for i in 1:k - pullback_i = Zygote.pullback(inv, A[:, :, i]) - - @test pullback_total[1][:, :, i] ≈ pullback_i[1] - @test in_diff[:, :, i] ≈ pullback_i[2](out_diff[:, :, i])[1] - end -end - -test22_inverse_pullback() diff --git a/test/quality/jet.jl b/test/quality/jet.jl index 001b08b8a..ad5934a0d 100644 --- a/test/quality/jet.jl +++ b/test/quality/jet.jl @@ -58,14 +58,6 @@ const A3{T} = Array{T, 3} (2, 3, 2), zeros(T, 2, 3, 2), zeros(T, 2, 4, 2), zeros(T, 3, 4, 2))) end - # src/kernels/identity_tensor.jl - for T in (Float16, Float32, Float64) - @test isempty(JET.get_reports(JET.report_opt(GML.assign_ones!, (A3{T},); - target_modules = TARGET))) - @test isempty(kernel_body_reports(GML.assign_ones_kernel!, (3, 2), - zeros(T, 3, 3, 2))) - end - # src/kernels/cayley.jl: `tensor_cayley`, its kernel and its pullback for T in (Float16, Float32, Float64) @test isempty(JET.get_reports(JET.report_opt(GML.tensor_cayley, (A3{T},); @@ -88,16 +80,8 @@ const A3{T} = Array{T, 3} @test isempty(kernel_body_reports(GML.da_kernel!, (3, 3, 2), zeros(3, 3, 2), Z, M, zeros(4, 4, 2), 3, 4)) - # src/kernels/inverses/: the symbolic inverses in test/kernels/tensor_inverse.jl + # src/kernels/inverses/tensor_mat_skew_sym_assign.jl for T in (Float16, Float32, Float64) - for (f, k, n) in ((GML.tensor_inverse2!, GML.inv22_kernel!, 2), - (GML.tensor_inverse3!, GML.inv33_kernel!, 3), - (GML.tensor_inverse4!, GML.inv44_kernel!, 4), - (GML.tensor_inverse5!, GML.inv55_kernel!, 5)) - @test isempty(JET.get_reports(JET.report_opt(f, (A3{T}, A3{T}); - target_modules = TARGET))) - @test isempty(kernel_body_reports(k, 2, zeros(T, n, n, 2), zeros(T, n, n, 2))) - end @test isempty(JET.get_reports(JET.report_opt(GML.tensor_mat_skew_sym_assign!, (A3{T}, A3{T}, Matrix{T}); target_modules = TARGET))) @test isempty(kernel_body_reports(GML.tensor_mat_skew_sym_assign_kernel!, diff --git a/test/runtests.jl b/test/runtests.jl index 09963a200..3b659f40a 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -16,11 +16,8 @@ if "core" in GROUPS @safetestset "The products of src/kernels/ on a JLArray" include("devices/jlarray.jl") @safetestset "Product gradients of structured and Stiefel weights" include("integration/product_gradients.jl") @safetestset "A Stiefel network trains on data of the other precision" include("integration/mixed_precision_training.jl") - @safetestset "Test parallel inverses" include("kernels/tensor_inverse.jl") - @safetestset "Test parallel Cayley" include("kernels/tensor_cayley.jl") @safetestset "tensor_cayley: skew inputs, the orthogonality bound and the pullback" include("kernels/cayley.jl") @safetestset "Volume-Preserving Transformer (skew-symmetric tests)" include("kernels/test_skew_map.jl") - @safetestset "Volume-Preserving Transformer (cayley-transform tests)" include("kernels/inverses/test_cayley_transforms.jl") @safetestset "Gradient Layer" include("layers/gradient_layer_tests.jl") @safetestset "Test tensor-slice consistency of sympnet layers" include("layers/sympnet_layers_test.jl") @safetestset "Manifold Neural Network Layers" include("layers/manifold_layers.jl")