diff --git a/CHANGELOG.md b/CHANGELOG.md index 18d295bb5..00130a1fd 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,16 @@ breaking release). ### Changed +- **`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, 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`. + + 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. + - `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 +118,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..5ff61182a 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` @@ -208,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` @@ -247,6 +211,33 @@ 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: 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 + 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 + 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 + ``` + ## Upstream ### K2 · Revise prints EMFILE errors in the test log 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/inverse_generator.jl b/scripts/utilities/inverse_generator.jl deleted file mode 100644 index 145fe96bf..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 `cpu_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 3c1c850f7..2758b4f70 100644 --- a/src/GeometricMachineLearning.jl +++ b/src/GeometricMachineLearning.jl @@ -127,14 +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/exponentials/tensor_exponential.jl") -include("kernels/inverses/cpu_inverse.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/cayley.jl") include("kernels/inverses/tensor_mat_skew_sym_assign.jl") include("kernels/kernel_ad_routines/seed.jl") diff --git a/src/kernels/cayley.jl b/src/kernels/cayley.jl new file mode 100644 index 000000000..0020914a9 --- /dev/null +++ b/src/kernels/cayley.jl @@ -0,0 +1,100 @@ +@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)``, 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) + 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. +# 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) + + @inbounds begin + for j in 1:n, i in 1:n + + δ = ifelse(i == j, one(eltype(C)), zero(eltype(C))) + A[i, j, k] = δ + C[i, j, k] + Λ[i, j, k] = δ - C[i, j, k] + end + + for col in 1:n + pivot = col + for r in (col + 1):n + abs(A[r, col, k]) > abs(A[pivot, col, k]) && (pivot = r) + end + + for j in 1:n + A[col, j, k], A[pivot, j, k] = A[pivot, j, k], A[col, j, k] + Λ[col, j, k], Λ[pivot, j, k] = Λ[pivot, j, k], Λ[col, j, k] + end + + d = A[col, col, k] + for j in 1:n + A[col, j, k] /= d + Λ[col, j, k] /= d + end + + for r in 1:n + r == col && continue + f = A[r, col, k] + for j in 1:n + A[r, j, k] -= f * A[col, j, k] + Λ[r, j, k] -= f * Λ[col, j, k] + end + end + end + end + + 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(Λ, Λ̄) + Y .+= Λ̄ + C̄ = tensor_tensor_transpose_mul(Y, Λ) + @. C̄ = -(Y + C̄) / T(2) + NoTangent(), C̄ + end + + Λ, tensor_cayley_pullback +end diff --git a/src/kernels/exponentials/tensor_exponential.jl b/src/kernels/exponentials/tensor_exponential.jl deleted file mode 100644 index 663dbfc92..000000000 --- a/src/kernels/exponentials/tensor_exponential.jl +++ /dev/null @@ -1,28 +0,0 @@ -# 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. - -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/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/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 d234d8d29..000000000 --- a/src/kernels/inverses/tensor_cayley.jl +++ /dev/null @@ -1,24 +0,0 @@ -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)) -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 2b7c4901f..5d299591c 100644 --- a/src/layers/volume_preserving_attention.jl +++ b/src/layers/volume_preserving_attention.jl @@ -3,9 +3,8 @@ 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. # Arguments @@ -20,10 +19,9 @@ 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 fast Cayley activation is using inverses that have been computed symbolically. +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. """ struct VolumePreservingAttention{M, N, ScalarProductType, SL} <: AbstractExplicitLayer{M, N} function VolumePreservingAttention(dim::Int, seq_length::Int = 0; skew_sym::Bool = false) @@ -37,31 +35,7 @@ end function orthonormal_activation_cayley(::VolumePreservingAttention, A::AbstractArray{ T, 3}) where {T} - cpu_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) + tensor_cayley(A) end function parameterlength(::VolumePreservingAttention{M, M, :skew_sym}) where {M} @@ -103,12 +77,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 ac08faf71..6cb059f1d 100644 --- a/test/devices/jlarray.jl +++ b/test/devices/jlarray.jl @@ -1,9 +1,9 @@ -# 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 @@ -12,6 +12,10 @@ using Zygote: gradient using Test import Random +# 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) @@ -59,3 +63,38 @@ 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 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:8 + + C = skew_batch(T, n, 4) + seed = rand(T, n, n, 4) + D = rand(T, n, n, 4) + + Λ, pullback = rrule(tensor_cayley, JLArray(C)) + _, C̄ = pullback(JLArray(seed)) + @test ondevice(Λ) && ondevice(C̄) + + 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 + # 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/devices/metal.jl b/test/devices/metal.jl index a12234279..df6b2b61d 100644 --- a/test/devices/metal.jl +++ b/test/devices/metal.jl @@ -13,12 +13,17 @@ # 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 Metal using Test import Random +# 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) Random.seed!(123) @@ -60,6 +65,30 @@ 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 + C = skew_batch(T, n, 8) + Λ = tensor_cayley(MtlArray(C)) + @test KernelAbstractions.get_backend(Λ) == backend + @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),)) + 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) <= + cayley_bound(C_layer) * maximum(abs, x) + end + end else @test_skip Metal.functional() end 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/kernels/cayley.jl b/test/kernels/cayley.jl new file mode 100644 index 000000000..48d4cbd3c --- /dev/null +++ b/test/kernels/cayley.jl @@ -0,0 +1,203 @@ +# `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, independent of the kernel. + +using GeometricMachineLearning +using GeometricMachineLearning: tensor_cayley, arbitrary_operand, + tensor_transpose_tensor_mul, _mul +using ChainRulesCore: AbstractZero, ZeroTangent, rrule +using LinearAlgebra: BLAS, I, opnorm +using Test +import Random + +# 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) + +# 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) + +# 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{T}, s) where {T} = C * T(s / largest_norm(C)) + +# (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)) +end + +# (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_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 +# `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 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} + T.(cayley_reference(BigFloat.(Ck))) +end + +# 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} + # 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) + worst = max(worst, orthogonality_error(cayley_reference(@view(C[:, :, k])))) + end + worst +end + +# 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))) + @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) + + # 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 + + # 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) + error_k = orthogonality_error(@view Λ[:, :, k]) + worst = max(worst, error_k / cayley_bound(T, @view(C[:, :, k]), n)) + end + @test worst <= 1 + end + + # 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) + 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, 3, 5) + + 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. 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 + # 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 + + cayley_rows(T, n, s) +end + +@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) + + 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 = cayley_bound(C) + @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/inverses/test_cayley_transforms.jl b/test/kernels/inverses/test_cayley_transforms.jl deleted file mode 100644 index 9c4b085a1..000000000 --- a/test/kernels/inverses/test_cayley_transforms.jl +++ /dev/null @@ -1,30 +0,0 @@ -using GeometricMachineLearning: tensor_cayley4, tensor_cayley3, cpu_tensor_cayley -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₃)) - 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 f4815a587..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, cpu_inverse, tensor_inverse2, - tensor_inverse3, tensor_inverse4, tensor_inverse5, - init_output, - _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,8 +10,8 @@ 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") function main(first_dim, second_dim, third_dim, third_tensor_dim) test_rrule(tensor_tensor_mul, rand(first_dim, second_dim, third_tensor_dim), @@ -27,8 +23,7 @@ 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_batch(Float64, first_dim, third_tensor_dim)) end const dim_range = 10 @@ -44,23 +39,13 @@ 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))), - (cpu_inverse, (invertible(T, 3, 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 @testset "a zero seed gives a zero cotangent for every argument" begin @@ -81,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 9a556a1d6..000000000 --- a/test/kernels/tensor_inverse.jl +++ /dev/null @@ -1,141 +0,0 @@ -using GeometricMachineLearning: tensor_inverse2, tensor_inverse3, tensor_inverse4, - tensor_inverse5, cpu_inverse -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() - -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/layers/volume_preserving_attention.jl b/test/layers/volume_preserving_attention.jl new file mode 100644 index 000000000..300b95a1f --- /dev/null +++ b/test/layers/volume_preserving_attention.jl @@ -0,0 +1,145 @@ +# `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 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; 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 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`. +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 +# `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, 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, N, T; skew_sym) + 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) + + # 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(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(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 "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, ps, x = case(0, seq_length, T; skew_sym) + + 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) + + # 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, 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..ad5934a0d 100644 --- a/test/quality/jet.jl +++ b/test/quality/jet.jl @@ -58,17 +58,19 @@ 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/cayley.jl: `tensor_cayley`, its kernel and its pullback for T in (Float16, Float32, Float64) - @test isempty(JET.get_reports(JET.report_opt(GML.assign_ones!, (A3{T},); + @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))) - @test isempty(kernel_body_reports(GML.assign_ones_kernel!, (3, 2), - zeros(T, 3, 3, 2))) 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) 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},); @@ -78,25 +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/: `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/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 ebc25815a..3b659f40a 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -16,10 +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") @@ -30,6 +28,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")