Products through GEMM: one _mul, slicing assign_q_and_p, zero and other-precision seeds in the kernel rrules (part M3)

KNOWN_ISSUES.md

@@ -111,22 +111,6 @@ What is known to be wrong in GeometricMachineLearning and is not fixed.
   deliberately: it is a restructuring with its own verification, not the comment repair that part
   was scoped to.
 
-### C15 · The `kernel_ad_routines` buffers are still zeroed where allocating would do.
-
-- **location:** `src/kernels/kernel_ad_routines/`
-- **kind:** defect
-- **found:** 2026-09-18
-- **evidence:**
-  `src/kernels/kernel_ad_routines/` allocates `dA = zero(A)`, `dS = KernelAbstractions.zeros(…)` and
-  the same shape in `tensor_mat_mul.jl`, `tensor_mat_skew_sym_assign.jl` and `vec_tensor_mul.jl`.
-  Their kernels sum into a local and assign the buffer element once, over an `ndrange` equal to the
-  buffer's size, so no element is read before it is written — the same argument that let the forward
-  wrappers move to `allocate` and `similar` under *Changed*.
-
-  This part was scoped to the forward path, so they were not changed with it. The backward pass is
-  where training spends its time, so the gain should be larger here than the forward figures, which
-  is also why it wants its own before-and-after measurement rather than being folded into that pass.
-
 ### C16 · One of `ExplicitImports`' seven checks is switched off in `test/quality/aqua.jl`, and two run with a list of names.
 
 - **location:** `test/quality/aqua.jl`
@@ -247,23 +231,6 @@ What is known to be wrong in GeometricMachineLearning and is not fixed.
   `test/reduced_system/reduced_system.jl:44–45` name `CHANGELOG.md` rather than `KNOWN_ISSUES.md`.
   The fix is a follow-up change to those comments.
 
-### K7 · A network with a Stiefel leaf raises in `mat_tensor_mul` when the data has another element type.
-
-- **location:** `src/kernels/mat_tensor_mul.jl`
-- **kind:** found late
-- **found:** 2026-10-12, by the critics of part M1a; the kernel code is older than that part
-- **evidence:**
-  A `Float32` `Chain(StiefelLayer(6, 3), Dense(3, 2, tanh))` trained with `train!` on `Float64`
-  data, and a `Float64` one on `Float32` data, raise `AssertionError: eltype(A) == eltype(B)` in
-  `mat_tensor_mul`, reached through `src/layers/manifold_layer.jl` under Zygote. A network of
-  `Dense` layers alone trains on mixed precisions.
-
-  ```julia
-  x = rand(Float32, 6, 12)
-  nn = NeuralNetwork(Chain(StiefelLayer(6, 3), Dense(3, 2, tanh)), Float64)
-  train!(nn, TrainingOptimizer(nn.params), DataLoader(x, x[1:2, :]), Batch(4), 2; loss = FeedForwardLoss())
-  ```
-
 ### K8 · A `DataLoader` of one sample cannot be batched.
 
 - **location:** `src/data_loader/batch.jl`
@@ -320,6 +287,45 @@ 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`
+- **kind:** defect
+- **found:** #345, review of 2026-10-15; the behaviour is the same on `origin/main` at `2a63f26a`
+- **evidence:**
+  `assign_q_and_p(x, N)` takes `x[1:N, …]` for `q` and the next `N` rows for `p`, and checks nothing
+  about `size(x, 1)`. An input with `size(x, 1) > 2N` loses its trailing rows silently, and
+  `PoissonTensor * v` reaches the same method.
+
+  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.
+
 ## Upstream
 
 ### K2 · Revise prints EMFILE errors in the test log

docs/src/arrays/tensors.md

@@ -5,13 +5,5 @@ We typically store training data as *tensors with three axes* in `GeometricMachi
 ## Library Functions
 
 ```@docs
-GeometricMachineLearning.tensor_mat_mul(::AbstractArray{<:Number, 3}, ::AbstractMatrix)
-GeometricMachineLearning.tensor_mat_mul!(::AbstractArray{<:Number, 3}, ::AbstractArray{<:Number, 3}, ::AbstractMatrix)
-GeometricMachineLearning.tensor_mat_mul!(::AbstractArray{<:Number, 3}, ::AbstractArray{<:Number, 3}, ::SymmetricMatrix)
-GeometricMachineLearning.mat_tensor_mul(::AbstractMatrix, ::AbstractArray{<:Number, 3})
-GeometricMachineLearning.mat_tensor_mul!(::AbstractArray{<:Number, 3}, ::AbstractMatrix, ::AbstractArray{<:Number, 3})
-GeometricMachineLearning.mat_tensor_mul!(::AbstractArray{<:Number, 3}, ::StrictlyLowerTriangular, ::AbstractArray{<:Number, 3})
-GeometricMachineLearning.mat_tensor_mul!(::AbstractArray{<:Number, 3}, ::StrictlyUpperTriangular, ::AbstractArray{<:Number, 3})
-GeometricMachineLearning.mat_tensor_mul!(::AbstractArray{<:Number, 3}, ::SkewSymMatrix, ::AbstractArray{<:Number, 3})
-GeometricMachineLearning.mat_tensor_mul!(::AbstractArray{<:Number, 3}, ::SymmetricMatrix, ::AbstractArray{<:Number, 3})
-```
+GeometricMachineLearning._mul
+```

docs/src/index.md

@@ -52,13 +52,13 @@ Many layers have been adapted in order to be used for problems in scientific mac
 
 `GeometricMachineLearning` allocates and computes through `KernelAbstractions.jl` [churavy2020kernel](@cite), so its layers and architectures are written against any backend that package supports: `CUDA.jl` [besard2018juliagpu](@cite), `AMDGPU.jl`, `Metal.jl` [besard2022metal](@cite) and `oneAPI.jl` [besard2022one](@cite).
 
-**Little of that is tested.** `test/devices/metal.jl` runs two checks on an Apple GPU, in the `metal` test group, which the `Metal` job and the `macOS-latest` job of CI run: the products of a `PoissonTensor` with a wrapped device array, and the `tensor_mat_mul!` kernel. A green test matrix says nothing about the rest of the GPU path. What is written below was measured by hand on an Apple M4 Max through `Metal.jl` 1.11.1, in `Float32`, because Apple GPUs have no `Float64` at all, and against `GeometricOptimizers` 0.9.0.
+**Little of that is tested.** `test/devices/metal.jl` runs two testsets on an Apple GPU, in the `metal` test group, which the `Metal` job and the `macOS-latest` job of CI run: the products of a `PoissonTensor` with a wrapped device array; and the products of `_mul` — a dense and a `SymmetricMatrix` weight from the left, a dense one from the right — together with the `tensor_tensor_mul` kernel. A green test matrix says nothing about the rest of the GPU path. What is written below was measured by hand on an Apple M4 Max through `Metal.jl` 1.11.1, in `Float32`, because Apple GPUs have no `Float64` at all, and against `GeometricOptimizers` 0.9.0.
 
 Everything tried on that device ran. Constructed, applied to both a matrix and a 3-tensor, and copied back with `map_to_cpu`: `GSympNet`, `LASympNet`, `StandardTransformerIntegrator`, `LinearSymplecticTransformer`, `SymplecticTransformer` at both its default `transformer_dim` and an upscaling one, `VolumePreservingFeedForward`, `VolumePreservingTransformer`, `SymplecticAutoencoder`, `PSDArch` and `Transformer(…; Stiefel = true)`. The product of a `StrictlyLowerTriangular` matrix with a 3-tensor runs. `train!` trains a `GSympNet` for two epochs with its data in `MtlArray`s. `ClassificationTransformer` was not tried.
 
 Two of those results rest on `GeometricOptimizers` rather than on anything here.
 
-- `VolumePreservingFeedForward` and `VolumePreservingTransformer` applied to a **matrix** multiply a `StrictlyLowerTriangular` or `StrictlyUpperTriangular` weight by that matrix. `GeometricOptimizers` supplies a `KernelAbstractions` kernel for that product, so the call does not fall through to a generic multiply that reads one entry at a time, which `GPUArraysCore` refuses.
+- `VolumePreservingFeedForward` and `VolumePreservingTransformer` make their `StrictlyLowerTriangular` or `StrictlyUpperTriangular` weight dense before they multiply, as its product with the identity on the weight's own device. `GeometricOptimizers` supplies a `KernelAbstractions` kernel for that product, so the call does not fall through to a generic multiply that reads one entry at a time, which `GPUArraysCore` refuses.
 - `StiefelLayer`, `GrassmannLayer` and `PSDLayer`, and so every architecture that holds one of them, orthonormalize their weight at **construction** through `GeometricOptimizers.orthonormal_columns`. That is CholeskyQR2 — matrix products and triangular solves only — so it runs wherever the draw was allocated. `LinearAlgebra.qr!` is a host factorization and `Metal.jl` implements no `qr` for an `MtlArray`, so a device needs the other one.
 
 ## Tutorials 

scripts/utilities/inverse_generator.jl

@@ -1,6 +1,8 @@
 # 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. That reproduces the committed files byte for byte.
+# 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
@@ -71,10 +73,7 @@ function wrapper_source(n)
         out = tensor_inverse$(n)(A)
 
         function tensor_inverse_pullback(out_diff)
-            out_diff = unthunk(out_diff)
-
-            NoTangent(),
-            - tensor_transpose_tensor_mul(out, tensor_tensor_mul(out_diff, tensor_transpose(out)))
+            _tensor_inverse_pullback(out, out_diff)
         end
         out, tensor_inverse_pullback
     end

scripts/utilities/product_timings.jl

@@ -0,0 +1,72 @@
+# The forward and the backward time of the products in `src/kernels/mul.jl` at 10×10×1000, BLAS
+# on one thread: a left product with each weight kind, a right product with a dense and a symmetric
+# weight, and the pullback of `tensor_mat_skew_sym_assign`.
+#
+# Each variant runs in five fresh processes. A process makes one warm-up call, then takes the
+# minimum of 100 timed calls of the forward product and of its backward pass, the
+# `Zygote.gradient` through a `NetworkParameters`, with the hook of NeuralNetworkParameters. The
+# table gives the median of the five minima in µs. `tensor_mat_skew_sym_assign` has a backward
+# column only: its row times the pullback alone.
+#
+# Run from the repository root, with the scripts environment; one process of one variant is the
+# second form, which the first starts:
+#   julia --startup-file=no --project=scripts scripts/utilities/product_timings.jl
+#   julia --startup-file=no --project=scripts scripts/utilities/product_timings.jl left:SymmetricMatrix Float32
+
+using GeometricMachineLearning
+using GeometricMachineLearning: _mul, tensor_mat_skew_sym_assign
+using ChainRulesCore: rrule
+using LinearAlgebra: BLAS
+using NeuralNetworkParameters: NetworkParameters
+using Printf
+import Random
+import Zygote
+
+const VARIANTS = ("left:Matrix", "left:SymmetricMatrix", "left:SkewSymMatrix",
+    "left:StrictlyLowerTriangular", "left:StrictlyUpperTriangular", "right:Matrix",
+    "right:SymmetricMatrix", "pullback:tensor_mat_skew_sym_assign")
+const PROCESSES = 5
+const CALLS = 100
+
+minimum_time(f) = (f(); minimum(_ -> @elapsed(f()), 1:CALLS) * 1e6)
+median(v) = sort(v)[(length(v) + 1) ÷ 2]
+
+function weight(kind, ::Type{T}) where {T}
+    kind == "Matrix" ? rand(T, 10, 10) :
+    rand(getfield(GeometricMachineLearning, Symbol(kind)){T}, 10)
+end
+
+# The forward and the backward minimum of one variant, in µs, in this process.
+function measure(variant, ::Type{T}) where {T}
+    BLAS.set_num_threads(1)
+    Random.seed!(0)
+    side, kind = split(variant, ':')
+    x = rand(T, 10, 10, 1000)
+    if side == "pullback"
+        pullback = rrule(tensor_mat_skew_sym_assign, x, rand(T, 10, 10))[2]
+        dB = rand(T, 10, 10, 1000)
+        return NaN, minimum_time(() -> pullback(dB))
+    end
+    product = side == "left" ? (W -> _mul(W, x)) : (W -> _mul(x, W))
+    ps = NetworkParameters((L1 = (W = weight(kind, T),),))
+    forward = minimum_time(() -> product(ps.L1.W))
+    backward = minimum_time(() -> Zygote.gradient(p -> sum(abs2, product(p.L1.W)), ps))
+    forward, backward
+end
+
+function one_process(variant, T)
+    cmd = `$(Base.julia_cmd()) --startup-file=no --project=$(Base.active_project()) $(@__FILE__) $variant $T`
+    parse.(Float64, split(readchomp(cmd)))
+end
+
+if length(ARGS) == 2
+    println(join(measure(ARGS[1], getfield(Base, Symbol(ARGS[2]))), " "))
+else
+    @printf("%-38s %-8s %12s %12s\n", "variant", "T", "forward µs", "backward µs")
+    for variant in VARIANTS, T in (Float32, Float64)
+
+        runs = [one_process(variant, T) for _ in 1:PROCESSES]
+        @printf("%-38s %-8s %12.1f %12.1f\n", variant, T,
+            median(first.(runs)), median(last.(runs)))
+    end
+end

src/GeometricMachineLearning.jl

@@ -123,33 +123,24 @@ export StiefelProjection
 export PoissonTensor
 
 include("kernels/assign_q_and_p.jl")
-include("kernels/tensor_mat_mul.jl")
 include("kernels/tensor_tensor_mul.jl")
 include("kernels/tensor_transpose_tensor_mul.jl")
 include("kernels/tensor_tensor_transpose_mul.jl")
-include("kernels/tensor_transpose_mat_mul.jl")
-include("kernels/mat_tensor_mul.jl")
-include("kernels/tensor_transpose.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/inverses/tensor_mat_skew_sym_assign.jl")
-include("kernels/vec_tensor_mul.jl")
 
-include("kernels/kernel_ad_routines/matrix_cotangent.jl")
-include("kernels/kernel_ad_routines/assign_q_and_p.jl")
-include("kernels/kernel_ad_routines/tensor_mat_mul.jl")
-include("kernels/kernel_ad_routines/mat_tensor_mul.jl")
+include("kernels/kernel_ad_routines/seed.jl")
 include("kernels/kernel_ad_routines/tensor_tensor_mul.jl")
-include("kernels/kernel_ad_routines/tensor_transpose_mat_mul.jl")
 include("kernels/kernel_ad_routines/tensor_transpose_tensor_mul.jl")
-include("kernels/kernel_ad_routines/tensor_transpose.jl")
 include("kernels/kernel_ad_routines/tensor_mat_skew_sym_assign.jl")
-include("kernels/kernel_ad_routines/vec_tensor_mul.jl")
 
 export MatrixSoftmax, VectorSoftmax
 include("activations/softmax.jl")

src/kernels/assign_q_and_p.jl

@@ -1,71 +1,7 @@
-"""
-This splits the input x into q and p part (works with kernels also on GPU). For now this is only used in the Sympnet.
-"""
-
-@kernel function assign_first_half!(q::AbstractVector, x::AbstractVector)
-    i = @index(Global)
-    q[i] = x[i]
-end
-
-@kernel function assign_second_half!(p::AbstractVector, x::AbstractVector, N::Integer)
-    i = @index(Global)
-    p[i] = x[i + N]
-end
-
-@kernel function assign_first_half!(q::AbstractMatrix, x::AbstractMatrix)
-    i, j = @index(Global, NTuple)
-    q[i, j] = x[i, j]
-end
-
-@kernel function assign_second_half!(p::AbstractMatrix, x::AbstractMatrix, N::Integer)
-    i, j = @index(Global, NTuple)
-    p[i, j] = x[i + N, j]
-end
-
-@kernel function assign_first_half!(q::AbstractArray{T, 3}, x::AbstractArray{
-        T, 3}) where {T}
-    i, j, k = @index(Global, NTuple)
-    q[i, j, k] = x[i, j, k]
-end
-
-@kernel function assign_second_half!(
-        p::AbstractArray{T, 3}, x::AbstractArray{
-            T, 3}, N::Integer) where {T}
-    i, j, k = @index(Global, NTuple)
-    p[i, j, k] = x[i + N, j, k]
-end
-
-# Allocate `q` and `p`, each with `N` as its first dimension -- half of the input's, which `N` also
-# states -- and return them as a `(q, p)` `NamedTuple`.
-function assign_q_and_p(x::AbstractVector, N::Int)
-    backend = networkbackend(x)
-    q = KernelAbstractions.allocate(backend, eltype(x), N)
-    p = KernelAbstractions.allocate(backend, eltype(x), N)
-    q_kernel! = assign_first_half!(backend)
-    p_kernel! = assign_second_half!(backend)
-    q_kernel!(q, x, ndrange = size(q))
-    p_kernel!(p, x, N, ndrange = size(p))
-    (q = q, p = p)
-end
-
-function assign_q_and_p(x::AbstractMatrix, N::Int)
-    backend = networkbackend(x)
-    q = KernelAbstractions.allocate(backend, eltype(x), N, size(x, 2))
-    p = KernelAbstractions.allocate(backend, eltype(x), N, size(x, 2))
-    q_kernel! = assign_first_half!(backend)
-    p_kernel! = assign_second_half!(backend)
-    q_kernel!(q, x, ndrange = size(q))
-    p_kernel!(p, x, N, ndrange = size(p))
-    (q = q, p = p)
-end
-
-function assign_q_and_p(x::AbstractArray{T, 3}, N::Int) where {T}
-    backend = networkbackend(x)
-    q = KernelAbstractions.allocate(backend, T, N, size(x, 2), size(x, 3))
-    p = KernelAbstractions.allocate(backend, T, N, size(x, 2), size(x, 3))
-    q_kernel! = assign_first_half!(backend)
-    p_kernel! = assign_second_half!(backend)
-    q_kernel!(q, x, ndrange = size(q))
-    p_kernel!(p, x, N, ndrange = size(p))
-    (q = q, p = p)
+# Split `x` along its first axis into the `(q, p)` `NamedTuple` of its first `N` rows and the `N`
+# rows after them. Rows from `2N + 1` on are ignored. Slicing works on every backend, and Zygote
+# differentiates it, so the split needs neither a kernel nor an `rrule`.
+function assign_q_and_p(x::AbstractArray, N::Integer)
+    rest = ntuple(Returns(:), ndims(x) - 1)
+    (q = x[1:N, rest...], p = x[(N + 1):(2N), rest...])
 end

src/kernels/inverses/cpu_inverse.jl

@@ -33,7 +33,8 @@ function ChainRulesCore.rrule(::typeof(cpu_inverse), A::AbstractArray)
     B = cpu_inverse(A)
 
     function cpu_inverse_pullback(dB)
-        dB = unthunk(dB)
+        dB = _seed(eltype(A), dB)
+        dB isa AbstractZero && return NoTangent(), ZeroTangent()
         dA = zero(dB)
         backend = networkbackend(dB)
 
@@ -45,9 +46,3 @@ function ChainRulesCore.rrule(::typeof(cpu_inverse), A::AbstractArray)
 
     return B, cpu_inverse_pullback
 end
-
-function cpu_tensor_cayley(A::AbstractArray)
-    one_A = init_output(A)
-
-    tensor_tensor_mul(one_A - A, cpu_inverse(one_A + A))
-end

src/kernels/inverses/inverse_2x2.jl

@@ -50,10 +50,7 @@ function ChainRulesCore.rrule(::typeof(tensor_inverse2), A::AT) where {
     out = tensor_inverse2(A)
 
     function tensor_inverse_pullback(out_diff)
-        out_diff = unthunk(out_diff)
-
-        NoTangent(),
-        - tensor_transpose_tensor_mul(out, tensor_tensor_mul(out_diff, tensor_transpose(out)))
+        _tensor_inverse_pullback(out, out_diff)
     end
     out, tensor_inverse_pullback
 end