diff --git a/CHANGELOG.md b/CHANGELOG.md index 1ac5f37a0..7dc1acf26 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -63,17 +63,27 @@ breaking release). - **This closes issues #319 and #343.** +- **The gradients of the structured kernel products are no longer doubled.** At the base `2a63f26a`, the structured kernel rrules returned a structural cotangent that NeuralNetworkParameters 0.5's `storage_gradient` doubled again. A `LinearLayerQ(4)` weight's gradient on a 4×1×7 tensor was 30.888, where the same data as a 4×7 matrix gave 15.444, and the finite-difference value is 15.444; `LinearSymplecticAttentionQ` gave 15.33 against 7.666, and `SymplecticAttentionQ` 0.9359 against 0.468. A `SymmetricMatrix` or `SkewSymMatrix` weight through `mat_tensor_mul` was off by a factor of two on the off-diagonal storage entries and on all of them, respectively. Every structured gradient now equals `central_fdm(5, 1)` on the storage, at `rtol = 1e-8` in Float64 and `√eps(Float32)` in Float32. The test is `test/integration/product_gradients.jl`. + +- **`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. + +- **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`. + ### Added - **`KernelAbstractions.get_backend` is defined for `PoissonTensor`.** GeometricOptimizers 0.9 asks every operand of a checked product for its backend, so `𝕁 * Y` with a `StiefelManifold` raised `ArgumentError: Implement KernelAbstractions.get_backend(::PoissonTensor...)` without it. -- **The test suite runs on an Apple GPU in its `metal` group.** `test/devices/metal.jl` checks the two device paths that were measured by hand and not tested before: the products of a `PoissonTensor` with a wrapped device array, `𝕁 * view(A, [1, 2, 3, 4], :)` and `𝕁 * B'`, which `ext/GPUArraysCoreExt.jl` carries; and the `tensor_mat_mul!` kernel on `MetalBackend()`. Each result must stay on the device and match the host in `Float32`, with scalar indexing off. `Pkg.test(test_args = ["metal"])` runs them without the subject tests, and a missing device records a skip, which happens inside a sandbox. Metal (1.10 or later) is a test dependency on every platform; it installs and precompiles on Linux and Windows, and nothing there loads it. On an Apple M4 Max with Metal 1.11.1 all five assertions pass. +- **The test suite runs on an Apple GPU in its `metal` group.** `test/devices/metal.jl` checks the two device paths that were measured by hand and not tested before: the products of a `PoissonTensor` with a wrapped device array, `𝕁 * view(A, [1, 2, 3, 4], :)` and `𝕁 * B'`, which `ext/GPUArraysCoreExt.jl` carries; and, in the testset "_mul and tensor_tensor_mul on Metal match the host", `_mul` with a dense and a `SymmetricMatrix` weight from the left, a dense one from the right, and `tensor_tensor_mul` on `MetalBackend()`. Each result must stay on the device and match the host in `Float32`, with scalar indexing off. `Pkg.test(test_args = ["metal"])` runs them without the subject tests, and a missing device records a skip, which happens inside a sandbox. Metal (1.10 or later) is a test dependency on every platform; it installs and precompiles on Linux and Windows, and nothing there loads it. On an Apple M4 Max with Metal 1.11.1 and Julia 1.13.1, both testsets pass in a cold process: 4 of 4 assertions in the `PoissonTensor` testset and 8 of 8 in the "_mul and tensor_tensor_mul on Metal match the host" testset. - **A `Metal` workflow runs the Metal tests on GitHub's `macos-15` runner**, for the `min` and `1` Julia versions. The runner's GPU is Apple's paravirtual device, which Metal.jl supports from 1.10 and on macOS 15 or later. The workflow includes a step "Metal is functional" before the tests that fails the job where the device is unavailable, so it cannot pass without having run the tests. The job is not a required check. However, the Metal tests also run in two required checks—CI.yml's `Julia min - macOS-latest - default` and `Julia 1 - macOS-latest - default`—which execute `test/runtests.jl` with no test arguments on Apple silicon, selecting the `metal` group. - **`test/quality/jet.jl` checks the kernel launchers with JET.** It runs in `core`, directly after `quality/aqua.jl`, and takes about 14 s. Each function that launches a kernel of `src/` on the `CPU()` backend, the pullbacks among them, and the `Batch` functor that `test/data_loader/batch_index_set.jl` measures with `@allocated` get one `JET.report_opt` line per element type that the tests use, with `target_modules = (GeometricMachineLearning,)`. Each `@kernel` method also gets a line that analyses its generated `cpu_` function at the launcher's context, because JET drops the reports of a kernel body when it analyses the launcher. 142 lines pass on Julia 1.11.9 (JET 0.9.20) and 1.13.1 (JET 0.12.2). Three lines report a runtime dispatch on both versions and are `@test_broken` with issue #325: `PoissonTensor(::CPU, ::Int, ::DataType)`, and the input-and-output method of `convert_input_and_batch_indices_to_array` at `Float32` and `Float64`. Where JET does not work, the file records one `@test_skip`. JET joins `test/Project.toml` with no `[compat]` bound. `KNOWN_ISSUES.md` K2 records the Revise `EMFILE` messages that JET brings into the test log, and K3 the dispatches that the lines do not see: an unstable argument of a kernel launch, and two barriers in the Poisson-tensor kernel. - **A test calls `iterate` on `(q,)` initial conditions.** No test reached the `NeuralNetworkIntegrator` `iterate` method that takes `ics::NamedTuple{(:q,), Tuple{AT}}`. `test/docstrings/layers_and_architectures.jl` runs the docstring's `ResNet` from `(q = [1, 1, 1],)` for four points and expects the same trajectory as the vector form, `[1 2 4 8; 1 3 9 27; 1 3 7 15]`. - **`NetworkLoss(arch)` is defined for a `NeuralNetworkIntegrator` `arch` and returns `FeedForwardLoss()`.** It is the default loss of that architecture: `train!` takes it for a `Batch{:FeedForward}` when no `loss` is given. A caller can name the default, or pass it explicitly, without relying on `train!`'s keyword default. +- **`Pkg.test` gains a `JLArray` device test in the `core` group.** `test/devices/jlarray.jl` runs `tensor_tensor_mul` and its pullback, `_mul` with each structured weight in the forward pass and in the storage gradient, and `assign_q_and_p` on a `JLArray` under `allowscalar(false)`. `JLArrays`, `GPUArraysCore` and `ChainRulesCore` join `test/Project.toml`, with `JLArrays = "0.3"`. `NNlib.batched_mul` is not used for `tensor_tensor_mul`, because on a backend with no NNlib extension, JLArray and Metal among them, it falls back to scalar indexing. + ### Changed - `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. @@ -86,6 +96,30 @@ 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. + + 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. + + | variant | T | forward | backward | + |:--|:--|--:|--:| + | left, Matrix | Float32 | 350.2 → 66.6 | 773.1 → 231.0 | + | left, Matrix | Float64 | 345.2 → 83.2 | 779.0 → 282.2 | + | left, SymmetricMatrix | Float32 | 602.4 → 70.1 | 2567.8 → 237.5 | + | left, SymmetricMatrix | Float64 | 592.3 → 85.1 | 2529.0 → 286.3 | + | left, SkewSymMatrix | Float32 | 534.0 → 71.7 | 4700.2 → 236.4 | + | left, SkewSymMatrix | Float64 | 528.6 → 84.3 | 4709.5 → 279.9 | + | left, StrictlyLowerTriangular | Float32 | 318.0 → 63.8 | 3421.3 → 237.6 | + | left, StrictlyLowerTriangular | Float64 | 312.8 → 84.0 | 3420.0 → 280.1 | + | left, StrictlyUpperTriangular | Float32 | 302.8 → 68.8 | 1462.0 → 229.3 | + | left, StrictlyUpperTriangular | Float64 | 303.9 → 90.3 | 1427.5 → 288.0 | + | right, Matrix | Float32 | 350.8 → 99.6 | 794.7 → 300.2 | + | right, Matrix | Float64 | 346.2 → 122.8 | 785.4 → 370.4 | + | right, SymmetricMatrix | Float32 | 534.0 → 86.9 | 7919.9 → 279.1 | + | right, SymmetricMatrix | Float64 | 538.3 → 121.5 | 8022.2 → 355.3 | + | `tensor_mat_skew_sym_assign` pullback (C15) | Float32 | — | 8365.5 → 8315.0 | + | `tensor_mat_skew_sym_assign` pullback (C15) | Float64 | — | 8387.2 → 8354.8 | + +- **The kernels that zero their output are gone, and `tensor_mat_skew_sym_assign`'s pullback allocates without zeroing.** After the deletions in the entry above, the only zeroed buffers left under `src/kernels/kernel_ad_routines/` were the two of `tensor_mat_skew_sym_assign`'s pullback. Both kernels assign every element once over `ndrange = size`, so both now use `similar` or `KernelAbstractions.allocate`, and no accumulating rule remains. This closes known issue C15, which leaves `KNOWN_ISSUES.md`. ### Fixed @@ -95,6 +129,7 @@ breaking release). - **The *GPU Support* section of `docs/src/index.md` names the two device paths the suite now tests**, where it said that nothing on the GPU path was tested. - **The manual's 27 TikZ figures and its logo come from [GeometricFigures](https://juliagni.github.io/GeometricFigures.jl/), and `docs/src/tikz/` with its 65 sources is gone.** Each page links the published light and dark SVG of a figure, and the 300-dpi PNG of the Grassmann sampling figure; `docs/make.jl` loads GeometricFigures' `figures.css` as a remote asset to show the one that the theme needs, and `extra_styles.css` loses its two `prefers-color-scheme` image rules. `make.jl` downloads the logo, and a failed download throws. Documenter's LaTeX writer does not fetch a remote image, so for `latex_output` `make.jl` downloads the light PDF of each figure that the `.tex` files name into `docs/build/` and puts the file name in place of the URL; `Latex.yml` and `docs/Makefile` no longer build the TikZ sources, and `docs/Makefile` drops its `latex_no_images`, `latex_no_pdf_no_images` and `html_no_images` targets: `latex`, `latex_no_pdf` and `html` run their recipes. `docs/check_references.jl` runs at the top of `make.jl`, before `makedocs`, and throws on an unresolved reference instead of calling `exit`; `test/quality/docs_make_checks_references.jl` checks that a bad `@ref` in a docstring stops `make.jl` there. `.github/workflows/Documenter.yml` is the canonical template again, without its TeX, TikZ and reference-check steps. +- **The manual documents `_mul` in place of the removed products.** `docs/src/arrays/tensors.md` lists `GeometricMachineLearning._mul` where it listed nine `@docs` lines for the old products, and its docstring carries the doctests of the old `mat_tensor_mul` and `tensor_mat_mul`, a left and a right product. The *GPU Support* section of `docs/src/index.md` is updated to match. ## [0.8.0] — 2026-09-21 diff --git a/KNOWN_ISSUES.md b/KNOWN_ISSUES.md index 64c528282..eb936965a 100644 --- a/KNOWN_ISSUES.md +++ b/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 diff --git a/docs/src/arrays/tensors.md b/docs/src/arrays/tensors.md index 0a8d65a9d..e3f3e53bc 100644 --- a/docs/src/arrays/tensors.md +++ b/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}) -``` \ No newline at end of file +GeometricMachineLearning._mul +``` diff --git a/docs/src/index.md b/docs/src/index.md index 9c8b9bdd5..1771763c5 100644 --- a/docs/src/index.md +++ b/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 diff --git a/scripts/utilities/inverse_generator.jl b/scripts/utilities/inverse_generator.jl index 055a7efb7..145fe96bf 100644 --- a/scripts/utilities/inverse_generator.jl +++ b/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 diff --git a/scripts/utilities/product_timings.jl b/scripts/utilities/product_timings.jl new file mode 100644 index 000000000..6efd3257a --- /dev/null +++ b/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 diff --git a/src/GeometricMachineLearning.jl b/src/GeometricMachineLearning.jl index 19aa4f1ba..3c1c850f7 100644 --- a/src/GeometricMachineLearning.jl +++ b/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") diff --git a/src/kernels/assign_q_and_p.jl b/src/kernels/assign_q_and_p.jl index ee674ffc2..e3a05577e 100644 --- a/src/kernels/assign_q_and_p.jl +++ b/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 diff --git a/src/kernels/inverses/cpu_inverse.jl b/src/kernels/inverses/cpu_inverse.jl index 4a173e34c..f21cdc478 100644 --- a/src/kernels/inverses/cpu_inverse.jl +++ b/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 diff --git a/src/kernels/inverses/inverse_2x2.jl b/src/kernels/inverses/inverse_2x2.jl index 83a979d82..0c80ecdcf 100644 --- a/src/kernels/inverses/inverse_2x2.jl +++ b/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 diff --git a/src/kernels/inverses/inverse_3x3.jl b/src/kernels/inverses/inverse_3x3.jl index f116e360f..f8afcfad7 100644 --- a/src/kernels/inverses/inverse_3x3.jl +++ b/src/kernels/inverses/inverse_3x3.jl @@ -89,10 +89,7 @@ function ChainRulesCore.rrule(::typeof(tensor_inverse3), A::AT) where { out = tensor_inverse3(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 diff --git a/src/kernels/inverses/inverse_4x4.jl b/src/kernels/inverses/inverse_4x4.jl index b322f6e74..a52c72a83 100644 --- a/src/kernels/inverses/inverse_4x4.jl +++ b/src/kernels/inverses/inverse_4x4.jl @@ -200,10 +200,7 @@ function ChainRulesCore.rrule(::typeof(tensor_inverse4), A::AT) where { out = tensor_inverse4(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 diff --git a/src/kernels/inverses/inverse_5x5.jl b/src/kernels/inverses/inverse_5x5.jl index dff844049..7686f7a49 100644 --- a/src/kernels/inverses/inverse_5x5.jl +++ b/src/kernels/inverses/inverse_5x5.jl @@ -478,10 +478,7 @@ function ChainRulesCore.rrule(::typeof(tensor_inverse5), A::AT) where { out = tensor_inverse5(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 diff --git a/src/kernels/inverses/tensor_inverse_pullback.jl b/src/kernels/inverses/tensor_inverse_pullback.jl new file mode 100644 index 000000000..1a0ebf1bd --- /dev/null +++ b/src/kernels/inverses/tensor_inverse_pullback.jl @@ -0,0 +1,10 @@ +# 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/kernels/kernel_ad_routines/assign_q_and_p.jl b/src/kernels/kernel_ad_routines/assign_q_and_p.jl deleted file mode 100644 index 2e5a4969a..000000000 --- a/src/kernels/kernel_ad_routines/assign_q_and_p.jl +++ /dev/null @@ -1,19 +0,0 @@ -""" -This implements the custom pullback for assign_q_and_p -""" - -# `assign_q_and_p` returns a `(q, p)` `NamedTuple`, so its cotangent is one too -- and either the -# tangent itself or its two components may arrive as a thunk. Splatting it straight into `vcat` -# without unthunking concatenates nothing: it builds a two-element `Vector{Thunk}`, which is a -# silently wrong gradient rather than an error. -_vcat_qp(qp̄) = vcat(unthunk(qp̄.q), unthunk(qp̄.p)) - -function ChainRulesCore.rrule(::typeof(assign_q_and_p), x::AbstractArray, N::Integer) - qp = assign_q_and_p(x, N) - function assign_q_and_p_pullback(qp_diff) - f̄ = NoTangent() - concat = @thunk _vcat_qp(unthunk(qp_diff)) - return f̄, concat, NoTangent() - end - return qp, assign_q_and_p_pullback -end diff --git a/src/kernels/kernel_ad_routines/mat_tensor_mul.jl b/src/kernels/kernel_ad_routines/mat_tensor_mul.jl deleted file mode 100644 index 34742f6a8..000000000 --- a/src/kernels/kernel_ad_routines/mat_tensor_mul.jl +++ /dev/null @@ -1,279 +0,0 @@ -""" -This implements the custom pullback tor mat_tensor_mul -""" - -#the @thunk macro means that the computation is only performed in case it is needed -function ChainRulesCore.rrule(::typeof(mat_tensor_mul), B::AbstractMatrix{T}, A::AbstractArray{ - T, 3}) where {T} - @assert axes(A, 1) == axes(B, 2) - C = mat_tensor_mul(B, A) - function mat_tensor_mul_pullback(C_diff) - f̄ = NoTangent() - B_diff = @thunk _matrix_cotangent(B, sum(tensor_tensor_transpose_mul(unthunk(C_diff), A), dims = 3)) - A_diff = @thunk mat_tensor_mul(B', unthunk(C_diff)) - return f̄, B_diff, A_diff - end - return C, mat_tensor_mul_pullback -end - -################################ lower triangular - -@kernel function lower_da_kernel!(dA::AT, S::AbstractVector{T}, dC::AT) where { - T, AT <: AbstractArray{T, 3}} - l, m, h = @index(Global, NTuple) - - temp = zero(T) - for i in (l + 1):size(dA, 1) - temp += S[(i - 2) * (i - 1) ÷ 2 + l] * dC[i, m, h] - end - - dA[l, m, h] = temp - - nothing -end - -@kernel function lower_ds_kernel!(dS::AbstractMatrix{T}, A::AT, dC::AT) where { - T, AT <: AbstractArray{T, 3}} - l, h = @index(Global, NTuple) - n = size(A, 1) - - temp = zero(T) - for j in axes(A, 2) - for i in 2:n - i_sum = (i - 2) * (i - 1) ÷ 2 - temp = (i_sum < l < (i_sum + i)) ? temp + A[l - i_sum, j, h] * dC[i, j, h] : - temp - end - end - - dS[l, h] = temp - - nothing -end - -function ChainRulesCore.rrule(::typeof(lo_mat_mul), S::AbstractVector{T}, - A::AbstractArray{T, 3}, n::Int) where {T} - C = lo_mat_mul(S, A, n) - function lo_mat_mul_pullback(dC) - dC = unthunk(dC) - f̄ = NoTangent() - backend = networkbackend(dC) - lower_da! = lower_da_kernel!(backend) - lower_ds! = lower_ds_kernel!(backend) - - dA = zero(A) - dS = KernelAbstractions.zeros(backend, T, length(S), size(dC, 3)) - - lower_da!(dA, S, dC, ndrange = size(dA)) - lower_ds!(dS, A, dC, ndrange = size(dS)) - - return f̄, reshape(sum(dS, dims = 2), length(S)), dA, NoTangent() - end - - return C, lo_mat_mul_pullback -end - -function ChainRulesCore.rrule(::typeof(mat_tensor_mul), B::StrictlyLowerTriangular{T}, - A::AbstractArray{T, 3}) where {T} - @assert size(A, 1) == B.n - C = mat_tensor_mul(B, A) - function lower_triangular_mul_pullback(dC) - dC = unthunk(dC) - f̄, dS, dA, _ = rrule(lo_mat_mul, B.S, A, B.n)[2](dC) - - return f̄, StrictlyLowerTriangular(dS, B.n), dA - end - return C, lower_triangular_mul_pullback -end - -################################ upper triangular - -@kernel function upper_da_kernel!(dA::AT, S::AbstractVector{T}, dC::AT) where { - T, AT <: AbstractArray{T, 3}} - l, m, h = @index(Global, NTuple) - - temp = zero(T) - for i in 1:(l - 1) - temp += S[(l - 2) * (l - 1) ÷ 2 + i] * dC[i, m, h] - end - - dA[l, m, h] = temp - - nothing -end - -@kernel function upper_ds_kernel!(dS::AbstractMatrix{T}, A::AT, dC::AT) where { - T, AT <: AbstractArray{T, 3}} - l, h = @index(Global, NTuple) - n = size(A, 1) - - temp = zero(T) - for k in 1:n - k_sum = (k - 2) * (k - 1) ÷ 2 - if k_sum < l < (k_sum + k) - for j in axes(A, 2) - temp += A[k, j, h] * dC[l - k_sum, j, h] - end - end - end - - dS[l, h] = temp - - nothing -end - -function ChainRulesCore.rrule(::typeof(up_mat_mul), S::AbstractVector{T}, - A::AbstractArray{T, 3}, n::Int) where {T} - C = up_mat_mul(S, A, n) - function up_mat_mul_pullback(dC) - dC = unthunk(dC) - f̄ = NoTangent() - backend = networkbackend(dC) - upper_da! = upper_da_kernel!(backend) - upper_ds! = upper_ds_kernel!(backend) - - dA = zero(A) - dS = KernelAbstractions.zeros(backend, T, length(S), size(dC, 3)) - - upper_da!(dA, S, dC, ndrange = size(dA)) - upper_ds!(dS, A, dC, ndrange = size(dS)) - - return f̄, reshape(sum(dS, dims = 2), length(S)), dA, NoTangent() - end - - return C, up_mat_mul_pullback -end - -function ChainRulesCore.rrule(::typeof(mat_tensor_mul), B::StrictlyUpperTriangular{T}, - A::AbstractArray{T, 3}) where {T} - @assert size(A, 1) == B.n - C = mat_tensor_mul(B, A) - function upper_triangular_mul_pullback(dC) - dC = unthunk(dC) - f̄, dS, dA, _ = rrule(up_mat_mul, B.S, A, B.n)[2](dC) - - return f̄, StrictlyUpperTriangular(dS, B.n), dA - end - return C, upper_triangular_mul_pullback -end - -################################ skew-symmetric - -function ChainRulesCore.rrule(::typeof(skew_mat_mul), S::AbstractVector{T}, - A::AbstractArray{T, 3}, n::Int) where {T} - C = skew_mat_mul(S, A, n) - function skew_mat_mul_pullback(dC) - dC = unthunk(dC) - f̄ = NoTangent() - backend = networkbackend(dC) - lower_da! = lower_da_kernel!(backend) - lower_ds! = lower_ds_kernel!(backend) - upper_da! = upper_da_kernel!(backend) - upper_ds! = upper_ds_kernel!(backend) - - dA_lower = zero(A) - dS_lower = KernelAbstractions.zeros(backend, T, length(S), size(dC, 3)) - dA_upper = zero(A) - dS_upper = KernelAbstractions.zeros(backend, T, length(S), size(dC, 3)) - - lower_da!(dA_lower, S, dC, ndrange = size(dA_lower)) - lower_ds!(dS_lower, A, dC, ndrange = size(dS_lower)) - upper_da!(dA_upper, S, dC, ndrange = size(dA_upper)) - upper_ds!(dS_upper, A, dC, ndrange = size(dS_upper)) - - return f̄, reshape(sum(dS_lower - dS_upper, dims = 2), length(S)), - dA_lower - dA_upper, NoTangent() - end - - return C, skew_mat_mul_pullback -end - -function ChainRulesCore.rrule(::typeof(mat_tensor_mul), B::SkewSymMatrix{T}, A::AbstractArray{ - T, 3}) where {T} - @assert size(A, 1) == B.n - C = mat_tensor_mul(B, A) - function skew_sym_mul_pullback(dC) - dC = unthunk(dC) - f̄, dS, dA, _ = rrule(skew_mat_mul, B.S, A, B.n)[2](dC) - - return f̄, SkewSymMatrix(dS, B.n), dA - end - return C, skew_sym_mul_pullback -end - -################################ symmetric - -@kernel function symmetric_da_kernel!(dA::AT, S::AbstractVector{T}, dC::AT) where { - T, AT <: AbstractArray{T, 3}} - l, m, h = @index(Global, NTuple) - temp = zero(T) - for i in l:size(dA, 1) - i_sum = (i - 1) * i ÷ 2 - temp += S[i_sum + l] * dC[i, m, h] - end - l_sum = (l - 1) * l ÷ 2 - for i in 1:(l - 1) - temp += S[l_sum + i] * dC[i, m, h] - end - dA[l, m, h] = temp - - nothing -end - -@kernel function symmetric_ds_kernel!(dS::AbstractMatrix{T}, A::AT, dC::AT) where { - T, AT <: AbstractArray{T, 3}} - l, h = @index(Global, NTuple) - temp = zero(T) - for i in axes(dC, 1) - sum_i = (i - 1) * i ÷ 2 - if sum_i < l - for j in axes(dC, 2) - if 1 ≤ (l - sum_i) < i - temp += A[l - sum_i, j, h] * dC[i, j, h] - temp += A[i, j, h] * dC[l - sum_i, j, h] - end - if l == (sum_i + i) - temp += A[l - sum_i, j, h] * dC[i, j, h] - end - end - end - end - dS[l, h] = temp - - nothing -end - -function ChainRulesCore.rrule(::typeof(symmetric_mat_mul), S::AbstractVector{T}, - A::AbstractArray{T, 3}, n::Int) where {T} - C = symmetric_mat_mul(S, A, n) - function symmetric_mat_mul_pullback(dC) - dC = unthunk(dC) - backend = networkbackend(dC) - symmetric_da! = symmetric_da_kernel!(backend) - symmetric_ds! = symmetric_ds_kernel!(backend) - - dA = zero(A) - dS = KernelAbstractions.zeros(backend, T, length(S), size(dC, 3)) - - symmetric_da!(dA, S, dC, ndrange = size(dA)) - symmetric_ds!(dS, A, dC, ndrange = size(dS)) - - NoTangent(), reshape(sum(dS, dims = 2), length(S)), dA, NoTangent() - end - - C, symmetric_mat_mul_pullback -end - -function ChainRulesCore.rrule(::typeof(mat_tensor_mul), B::SymmetricMatrix{T}, A::AbstractArray{ - T, 3}) where {T} - @assert size(A, 1) == B.n - C = mat_tensor_mul(B, A) - function symmetric_mul_pullback(dC) - dC = unthunk(dC) - f̄, dS, dA, _ = rrule(symmetric_mat_mul, B.S, A, B.n)[2](dC) - - return f̄, SymmetricMatrix(dS, B.n), dA - end - - return C, symmetric_mul_pullback -end diff --git a/src/kernels/kernel_ad_routines/matrix_cotangent.jl b/src/kernels/kernel_ad_routines/matrix_cotangent.jl deleted file mode 100644 index f5a7bec78..000000000 --- a/src/kernels/kernel_ad_routines/matrix_cotangent.jl +++ /dev/null @@ -1,19 +0,0 @@ -# `_matrix_cotangent(B, dB)` gives the cotangent `dB` of a matrix argument `B` the shape *and* the -# array type of `B`. It is a plain comment rather than a docstring on purpose: `docs/make.jl` runs -# Documenter with the default `checkdocs = :all`, so every docstring in the module has to be placed -# in the manual. -# -# Two things need undoing. The kernel rules build a matrix argument's cotangent by contracting the -# third axis with `sum(_, dims = 3)`, which leaves a trailing singleton axis; a cotangent has to -# have the shape of the primal it belongs to, so that axis comes off. -# -# The array type matters for `Adjoint`, which is what `B` is wherever a layer multiplies by a -# transposed weight -- `MultiHeadAttention` computes `mat_tensor_mul(ps.PQ[key]', x)`. ChainRules' -# `rrule` for `adjoint` unwraps an `Adjoint` cotangent to its dense parent but leaves a dense one -# wrapped, so returning a plain `Matrix` here hands the caller an `Adjoint{T, Matrix{T}}` as the -# gradient of a `Matrix{T}` parameter. `ProjectTo` does not help: for a real `Adjoint` it is a -# `ProjectTo{AbstractArray}`, which keeps whatever wrapper the tangent arrived with. -_matrix_cotangent(::AbstractMatrix, dB::AbstractArray{<:Number, 3}) = dropdims(dB; dims = 3) -function _matrix_cotangent(::Adjoint, dB::AbstractArray{<:Number, 3}) - collect(dropdims(dB; dims = 3)')' -end diff --git a/src/kernels/kernel_ad_routines/seed.jl b/src/kernels/kernel_ad_routines/seed.jl new file mode 100644 index 000000000..81fe9529c --- /dev/null +++ b/src/kernels/kernel_ad_routines/seed.jl @@ -0,0 +1,6 @@ +# The seed `ȳ` of a kernel rule's pullback, for a primal of element type `T`: unthunked, and +# converted to `T`, because the kernels take one element type for all their arrays. An +# `AbstractZero` passes through, and every rule answers it with zeros before a kernel runs. +_seed(::Type{T}, ȳ) where {T} = _seed_eltype(T, unthunk(ȳ)) +_seed_eltype(::Type, ȳ::AbstractZero) = ȳ +_seed_eltype(::Type{T}, ȳ::AbstractArray) where {T} = convert(AbstractArray{T}, ȳ) diff --git a/src/kernels/kernel_ad_routines/tensor_mat_mul.jl b/src/kernels/kernel_ad_routines/tensor_mat_mul.jl deleted file mode 100644 index 791f86733..000000000 --- a/src/kernels/kernel_ad_routines/tensor_mat_mul.jl +++ /dev/null @@ -1,91 +0,0 @@ -""" -This implements the custom pullback for tensor_mat_mul -""" - -#the @thunk macro means that the computation is only performed in case it is needed -function ChainRulesCore.rrule(::typeof(tensor_mat_mul), A::AbstractArray{T, 3}, B::AbstractMatrix{T}) where {T} - @assert axes(A, 2) == axes(B, 1) - C = tensor_mat_mul(A, B) - function tensor_mat_mul_pullback(C_diff) - f̄ = NoTangent() - A_diff = @thunk tensor_mat_mul(unthunk(C_diff), B') - B_diff = @thunk _matrix_cotangent(B, sum(tensor_transpose_tensor_mul(A, unthunk(C_diff)), dims = 3)) - return f̄, A_diff, B_diff - end - return C, tensor_mat_mul_pullback -end - -############### Symmetric (right mul) - -@kernel function symmetric_right_da_kernel!( - dA::AT, S::AbstractVector{T}, dC::AT) where {T, AT <: AbstractArray{T, 3}} - l, m, h = @index(Global, NTuple) - - temp = zero(T) - - for j in 1:m - temp += S[(m - 1) * m ÷ 2 + j] * dC[l, j, h] - end - for j in (m + 1):size(dA, 2) - temp += S[(j - 1) * j ÷ 2 + m] * dC[l, j, h] - end - - dA[l, m, h] = temp - - nothing -end - -@kernel function symmetric_right_ds_kernel!( - dS::AbstractMatrix{T}, B::AT, dC::AT) where {T, AT <: AbstractArray{T, 3}} - l, h = @index(Global, NTuple) - temp = zero(T) - - for i in axes(dC, 1) - for k in axes(dC, 2) - sum_k = (k - 1) * k ÷ 2 - temp += 1 ≤ l - sum_k ≤ k ? B[i, k, h] * dC[i, l - sum_k, h] : zero(T) - end - for j in axes(dC, 2) - sum_j = (j - 1) * j ÷ 2 - temp += 1 ≤ l - sum_j < j ? B[i, l - sum_j, h] * dC[i, j, h] : zero(T) - end - end - - dS[l, h] = temp - - nothing -end - -function ChainRulesCore.rrule(::typeof(symmetric_mat_right_mul), A::AbstractArray{T, 3}, - S::AbstractVector{T}, n::Int) where {T} - C = symmetric_mat_right_mul(A, S, n) - function symmetric_mat_mul_pullback(dC) - dC = unthunk(dC) - backend = networkbackend(dC) - symmetric_right_da! = symmetric_right_da_kernel!(backend) - symmetric_right_ds! = symmetric_right_ds_kernel!(backend) - - dA = zero(A) - dS = KernelAbstractions.zeros(backend, T, length(S), size(dC, 3)) - - symmetric_right_da!(dA, S, dC, ndrange = size(dA)) - symmetric_right_ds!(dS, A, dC, ndrange = size(dS)) - - NoTangent(), dA, reshape(sum(dS, dims = 2), length(S)), NoTangent() - end - - C, symmetric_mat_mul_pullback -end - -function ChainRulesCore.rrule(::typeof(tensor_mat_mul), A::AbstractArray{T, 3}, B::SymmetricMatrix{T}) where {T} - @assert size(A, 2) == B.n - C = tensor_mat_mul(A, B) - function symmetric_right_mul_pullback(dC) - dC = unthunk(dC) - f̄, dA, dS, _ = rrule(symmetric_mat_right_mul, A, B.S, B.n)[2](dC) - - return f̄, dA, SymmetricMatrix(dS, B.n) - end - - return C, symmetric_right_mul_pullback -end diff --git a/src/kernels/kernel_ad_routines/tensor_mat_skew_sym_assign.jl b/src/kernels/kernel_ad_routines/tensor_mat_skew_sym_assign.jl index b59234d96..303447c25 100644 --- a/src/kernels/kernel_ad_routines/tensor_mat_skew_sym_assign.jl +++ b/src/kernels/kernel_ad_routines/tensor_mat_skew_sym_assign.jl @@ -38,14 +38,16 @@ function ChainRulesCore.rrule(::typeof(tensor_mat_skew_sym_assign), @assert size(A, 1) == size(Z, 1) B = tensor_mat_skew_sym_assign(Z, A) function tensor_mat_skew_sym_assign_pullback(dB) - dB = unthunk(dB) + dB = _seed(T, dB) + dB isa AbstractZero && return NoTangent(), ZeroTangent(), ZeroTangent() f̄ = NoTangent() backend = networkbackend(dB) dz! = dz_kernel!(backend) da! = da_kernel!(backend) - dZ = zero(Z) - dA = KernelAbstractions.zeros(backend, T, size(A)..., size(dB, 3)) + # Both kernels assign every element of their buffer once, so neither buffer is zeroed. + dZ = similar(Z) + dA = KernelAbstractions.allocate(backend, T, size(A)..., size(dB, 3)) sys_dim, seq_length, _ = size(Z) diff --git a/src/kernels/kernel_ad_routines/tensor_tensor_mul.jl b/src/kernels/kernel_ad_routines/tensor_tensor_mul.jl index f86522b60..4a0ceb6c5 100644 --- a/src/kernels/kernel_ad_routines/tensor_tensor_mul.jl +++ b/src/kernels/kernel_ad_routines/tensor_tensor_mul.jl @@ -1,17 +1,11 @@ -""" -This implements the custom pullback tor tensor_tensor_mul -""" - -#the @thunk macro means that the computation is only performed in case it is needed function ChainRulesCore.rrule(::typeof(tensor_tensor_mul), A::AbstractArray{T, 3}, B::AbstractArray{T, 3}) where {T} @assert axes(A, 2) == axes(B, 1) C = tensor_tensor_mul(A, B) - function tensor_tensor_mul_pullback(C_diff) - f̄ = NoTangent() - A_diff = @thunk tensor_tensor_transpose_mul(unthunk(C_diff), B) - B_diff = @thunk tensor_transpose_tensor_mul(A, unthunk(C_diff)) - return f̄, A_diff, B_diff + function tensor_tensor_mul_pullback(C̄) + C̄ = _seed(T, C̄) + C̄ isa AbstractZero && return NoTangent(), ZeroTangent(), ZeroTangent() + NoTangent(), tensor_tensor_transpose_mul(C̄, B), tensor_transpose_tensor_mul(A, C̄) end - return C, tensor_tensor_mul_pullback + C, tensor_tensor_mul_pullback end diff --git a/src/kernels/kernel_ad_routines/tensor_transpose.jl b/src/kernels/kernel_ad_routines/tensor_transpose.jl deleted file mode 100644 index ec1e276be..000000000 --- a/src/kernels/kernel_ad_routines/tensor_transpose.jl +++ /dev/null @@ -1,14 +0,0 @@ -""" -This implements the custom pullback tor tensor_transpose -""" - -#the @thunk macro means that the computation is only performed in case it is needed -function ChainRulesCore.rrule(::typeof(tensor_transpose), A::AbstractArray{T, 3}) where {T} - C = tensor_transpose(A) - function tensor_transpose_pullback(C_diff) - f̄ = NoTangent() - A_diff = @thunk tensor_transpose(unthunk(C_diff)) - return f̄, A_diff - end - return C, tensor_transpose_pullback -end diff --git a/src/kernels/kernel_ad_routines/tensor_transpose_mat_mul.jl b/src/kernels/kernel_ad_routines/tensor_transpose_mat_mul.jl deleted file mode 100644 index 11041c856..000000000 --- a/src/kernels/kernel_ad_routines/tensor_transpose_mat_mul.jl +++ /dev/null @@ -1,15 +0,0 @@ -# This implements the custom pullback for tensor_transpose_mat_mul -function ChainRulesCore.rrule(::typeof(tensor_transpose_mat_mul), - A::AbstractArray{T, 3}, B::AbstractMatrix{T}) where {T} - @assert axes(A, 1) == axes(B, 1) - C = tensor_transpose_mat_mul(A, B) - function tensor_transpose_mat_mul_pullback(C_diff) - f̄ = NoTangent() - A_diff = @thunk mat_tensor_transpose_mul(B, unthunk(C_diff)) - B_diff = @thunk _matrix_cotangent(B, sum(tensor_tensor_mul(A, unthunk(C_diff)), dims = 3)) - return f̄, A_diff, B_diff - end - return C, tensor_transpose_mat_mul_pullback -end - -mat_tensor_transpose_mul(B, C) = mat_tensor_mul(B, tensor_transpose(C)) diff --git a/src/kernels/kernel_ad_routines/tensor_transpose_tensor_mul.jl b/src/kernels/kernel_ad_routines/tensor_transpose_tensor_mul.jl index 4c5137bdd..fe685ea68 100644 --- a/src/kernels/kernel_ad_routines/tensor_transpose_tensor_mul.jl +++ b/src/kernels/kernel_ad_routines/tensor_transpose_tensor_mul.jl @@ -1,17 +1,11 @@ -""" -This implements the custom pullback tor tensor_transpose_tensor_mul -""" - -#the @thunk macro means that the computation is only performed in case it is needed function ChainRulesCore.rrule(::typeof(tensor_transpose_tensor_mul), A::AbstractArray{T, 3}, B::AbstractArray{T, 3}) where {T} @assert axes(A, 1) == axes(B, 1) C = tensor_transpose_tensor_mul(A, B) - function tensor_transpose_tensor_mul_pullback(C_diff) - f̄ = NoTangent() - A_diff = @thunk tensor_tensor_transpose_mul(B, unthunk(C_diff)) - B_diff = @thunk tensor_tensor_mul(A, unthunk(C_diff)) - return f̄, A_diff, B_diff + function tensor_transpose_tensor_mul_pullback(C̄) + C̄ = _seed(T, C̄) + C̄ isa AbstractZero && return NoTangent(), ZeroTangent(), ZeroTangent() + NoTangent(), tensor_tensor_transpose_mul(B, C̄), tensor_tensor_mul(A, C̄) end - return C, tensor_transpose_tensor_mul_pullback + C, tensor_transpose_tensor_mul_pullback end diff --git a/src/kernels/kernel_ad_routines/vec_tensor_mul.jl b/src/kernels/kernel_ad_routines/vec_tensor_mul.jl deleted file mode 100644 index e373a6114..000000000 --- a/src/kernels/kernel_ad_routines/vec_tensor_mul.jl +++ /dev/null @@ -1,35 +0,0 @@ -function ChainRulesCore.rrule(::typeof(vec_tensor_mul), a::AbstractVector{T}, x::AbstractArray{ - T, 3}) where {T} - b = vec_tensor_mul(a, x) - function vec_tensor_mul_pullback(b_diff) - a_diff = @thunk tensor_scalar_product(x, unthunk(b_diff)) - x_diff = @thunk vec_tensor_mul(a, unthunk(b_diff)) - NoTangent(), a_diff, x_diff - end - b, vec_tensor_mul_pullback -end - -# This is computing the sum of scalar products for two tensors -@kernel function tensor_scalar_product_kernel!( - a_diff::AbstractVector{T}, x::AbstractArray{T, 3}, - b_diff::AbstractArray{T, 3}, range_2, range_3) where {T} - i = @index(Global) - a_val = zero(T) - for j in 1:range_2 - for k in 1:range_3 - a_val += x[i, j, k]*b_diff[i, j, k] - end - end - a_diff[i] = a_val -end - -function tensor_scalar_product(x::AbstractArray{T, 3}, b_diff::AbstractArray{ - T, 3}) where {T} - a_size = size(x, 1) - backend = networkbackend(x) - a_diff = KernelAbstractions.zeros(backend, T, a_size) - tensor_scalar_product! = tensor_scalar_product_kernel!(backend) - tensor_scalar_product!( - a_diff, x, b_diff, size(x, 2), size(x, 3), ndrange = size(a_diff)) - a_diff -end diff --git a/src/kernels/mat_tensor_mul.jl b/src/kernels/mat_tensor_mul.jl deleted file mode 100644 index 3e74ca9a3..000000000 --- a/src/kernels/mat_tensor_mul.jl +++ /dev/null @@ -1,262 +0,0 @@ -# `C[:, :, k] = A * B[:, :, k]`: a matrix times every slice of a tensor, one thread per output entry. -@kernel function mat_tensor_mul_kernel!(C, A, B) - i, j, k = @index(Global, NTuple) - - # creating a temporary sum variable for matrix multiplication - tmp_sum = zero(eltype(C)) - for l in axes(A, 2) - tmp_sum += A[i, l] * B[l, j, k] - end - - C[i, j, k] = tmp_sum - - nothing -end - -@doc raw""" - mat_tensor_mul!(C, A, B) - -Multiply the matrix `A` onto the tensor `B` from the left and store the result in `C`. - -The function [`mat_tensor_mul`](@ref) calls `mat_tensor_mul!` internally. -""" -function mat_tensor_mul!(C::AbstractArray{<:Number, 3}, A::AbstractMatrix, B::AbstractArray{ - <:Number, 3}) - @assert eltype(C) == eltype(A) == eltype(B) - @assert size(A)[2] == size(B)[1] - - backend = networkbackend(A) - kernel! = mat_tensor_mul_kernel!(backend) - kernel!(C, A, B, ndrange = size(C)) -end - -@doc raw""" - mat_tensor_mul(A, B) - -Multipliy the matrix `A` onto the tensor `B` from the left. - -Internally this calls the inplace version [`mat_tensor_mul!`](@ref). - -# Examples - -```jldoctest -using GeometricMachineLearning: mat_tensor_mul - -B = [1 1 1; 1 1 1; 1 1 1;;; 2 2 2; 2 2 2; 2 2 2] -A = [3 0 0; 0 2 0; 0 0 1] - -mat_tensor_mul(A, B) - -# output - -3×3×2 Array{Int64, 3}: -[:, :, 1] = - 3 3 3 - 2 2 2 - 1 1 1 - -[:, :, 2] = - 6 6 6 - 4 4 4 - 2 2 2 -``` -""" -function mat_tensor_mul(A::AbstractMatrix, B::AbstractArray{<:Number, 3}) - @assert eltype(A) == eltype(B) - T = eltype(A) - sizeA = size(A) - sizeB = size(B) - backend = networkbackend(A) - C = KernelAbstractions.allocate(backend, T, sizeA[1], sizeB[2], sizeB[3]) - mat_tensor_mul!(C, A, B) - C -end - -################### SymmetricMatrix - -@kernel function symmetric_mat_mul_kernel!(C::AbstractArray{T, 3}, S::AbstractVector{T}, - B::AbstractArray{T, 3}, n::Int) where {T} - i, j, l = @index(Global, NTuple) - - tmp_sum = zero(T) - for k in 1:i - tmp_sum += S[((i - 1) * i) ÷ 2 + k] * B[k, j, l] - end - for k in (i + 1):n - tmp_sum += S[((k - 1) * k) ÷ 2 + i] * B[k, j, l] - end - C[i, j, l] = tmp_sum -end - -function symmetric_mat_mul!(C::AbstractArray{T, 3}, S::AbstractVector{T}, - B::AbstractArray{T, 3}, n::Int) where {T} - backend = networkbackend(C) - - symmetric_mat_mul_k! = symmetric_mat_mul_kernel!(backend) - symmetric_mat_mul_k!(C, S, B, n, ndrange = size(C)) -end - -function symmetric_mat_mul(S::AbstractVector{T}, B::AbstractArray{T, 3}, n::Int) where {T} - C = similar(B) - - symmetric_mat_mul!(C, S, B, n) - - C -end - -@doc raw""" - mat_tensor_mul!(C, A::SymmetricMatrix, B) - -Multiply the symmetric matrix `A` onto the tensor `B` from the left and store the result in `C`. - -This performs an efficient multiplication based on the special structure of the symmetric matrix `A`. -""" -function mat_tensor_mul!(C::AbstractArray{<:Number, 3}, A::SymmetricMatrix, B::AbstractArray{ - <:Number, 3}) - @assert eltype(C) == eltype(A) == eltype(B) - @assert A.n == size(C, 1) == size(B, 1) - - symmetric_mat_mul!(C, A.S, B, A.n) -end - -########################### StrictlyLowerTriangular - -@kernel function lo_mul_kernel!(C::AbstractArray{T, 3}, S::AbstractVector{T}, - B::AbstractArray{T, 3}, ::Int) where {T} - i, j, l = @index(Global, NTuple) - - tmp_sum = zero(T) - for k in 1:(i - 1) - tmp_sum += S[(i - 2) * (i - 1) ÷ 2 + k] * B[k, j, l] - end - C[i, j, l] = tmp_sum - - nothing -end - -function lo_mat_mul!( - C::AbstractArray{T, 3}, S::AbstractVector{T}, B::AbstractArray{ - T, 3}, n::Int) where {T} - backend = networkbackend(C) - - lo_mat_mul_k! = lo_mul_kernel!(backend) - lo_mat_mul_k!(C, S, B, n, ndrange = size(C)) -end - -function lo_mat_mul(S::AbstractVector{T}, B::AbstractArray{T, 3}, n::Int) where {T} - C = similar(B) - - lo_mat_mul!(C, S, B, n) - - C -end - -@doc raw""" - mat_tensor_mul!(C, A::StrictlyLowerTriangular, B) - -Multiply the strictly lower-triangular matrix `A` onto the tensor `B` from the left and store the result in `C`. - -This performs an efficient multiplication based on the special structure of the strictly lower-triangular matrix `A`. -""" -function mat_tensor_mul!(C::AbstractArray{<:Number, 3}, A::StrictlyLowerTriangular, - B::AbstractArray{<:Number, 3}) - @assert eltype(C) == eltype(A) == eltype(B) - @assert A.n == size(C, 1) == size(B, 1) - - lo_mat_mul!(C, A.S, B, A.n) -end - -####################### StrictlyUpperTriangular - -@kernel function up_mul_kernel!(C::AbstractArray{T, 3}, S::AbstractVector{T}, - B::AbstractArray{T, 3}, n::Int) where {T} - i, j, l = @index(Global, NTuple) - - tmp_sum = zero(T) - for k in (i + 1):n - tmp_sum += S[(k - 2) * (k - 1) ÷ 2 + i] * B[k, j, l] - end - C[i, j, l] = tmp_sum - - nothing -end - -function up_mat_mul!( - C::AbstractArray{T, 3}, S::AbstractVector{T}, B::AbstractArray{ - T, 3}, n::Int) where {T} - backend = networkbackend(C) - - up_mat_mul_k! = up_mul_kernel!(backend) - up_mat_mul_k!(C, S, B, n, ndrange = size(C)) -end - -function up_mat_mul(S::AbstractVector{T}, B::AbstractArray{T, 3}, n::Int) where {T} - C = similar(B) - - up_mat_mul!(C, S, B, n) - - C -end - -@doc raw""" - mat_tensor_mul!(C, A::StrictlyUpperTriangular, B) - -Multiply the strictly upper-triangular matrix `A` onto the tensor `B` from the left and store the result in `C`. - -This performs an efficient multiplication based on the special structure of the strictly upper-triangular matrix `A`. -""" -function mat_tensor_mul!(C::AbstractArray{<:Number, 3}, A::StrictlyUpperTriangular, - B::AbstractArray{<:Number, 3}) - @assert eltype(C) == eltype(A) == eltype(B) - @assert A.n == size(C, 1) == size(B, 1) - - up_mat_mul!(C, A.S, B, A.n) -end - -################## SkewSymMatrix - -@kernel function skew_mat_mul_kernel!(C::AbstractArray{T, 3}, S::AbstractVector{T}, - B::AbstractArray{T, 3}, n::Int) where {T} - i, j, l = @index(Global, NTuple) - - tmp_sum = zero(T) - for k in 1:(i - 1) - tmp_sum += S[(i - 2) * (i - 1) ÷ 2 + k] * B[k, j, l] - end - for k in (i + 1):n - tmp_sum += -S[(k - 2) * (k - 1) ÷ 2 + i] * B[k, j, l] - end - C[i, j, l] = tmp_sum -end - -function skew_mat_mul!( - C::AbstractArray{T, 3}, S::AbstractVector{T}, B::AbstractArray{ - T, 3}, n::Int) where {T} - backend = networkbackend(C) - - skew_mat_mul_k! = skew_mat_mul_kernel!(backend) - skew_mat_mul_k!(C, S, B, n, ndrange = size(C)) -end - -function skew_mat_mul(S::AbstractVector{T}, B::AbstractArray{T, 3}, n::Int) where {T} - C = similar(B) - - skew_mat_mul!(C, S, B, n) - - C -end - -@doc raw""" - mat_tensor_mul!(C, A::SkewSymMatrix, B) - -Multiply skew-symmetric the matrix `A` onto the tensor `B` from the left and store the result in `C`. - -This performs an efficient multiplication based on the special structure of the skew-symmetric matrix `A`. -""" -function mat_tensor_mul!(C::AbstractArray{<:Number, 3}, A::SkewSymMatrix, B::AbstractArray{ - <:Number, 3}) - @assert eltype(C) == eltype(A) == eltype(B) - @assert A.n == size(C, 1) == size(B, 1) - - skew_mat_mul!(C, A.S, B, A.n) -end diff --git a/src/kernels/mul.jl b/src/kernels/mul.jl new file mode 100644 index 000000000..bf3fe74ff --- /dev/null +++ b/src/kernels/mul.jl @@ -0,0 +1,112 @@ +@doc raw""" + _mul(W::AbstractMatrix, x::AbstractArray) + _mul(x::AbstractArray{<:Number, 3}, W::AbstractMatrix) + _mul(x::AbstractArray{<:Number, 3}, y::AbstractArray{<:Number, 3}) + +Multiply a matrix and a tensor, or two tensors, slice by slice along the third axis. + +- The **left product** multiplies `W` onto `x` from the left. For a vector or a matrix `x` it is + `W * x`. For a tensor it is one GEMM: `x` is reshaped into a matrix with `size(x, 1)` rows, so + that `_mul(W, x)[:, :, k] == W * x[:, :, k]`. +- The **right product** multiplies `W` onto every slice of the tensor `x` from the right, + `_mul(x, W)[:, :, k] == x[:, :, k] * W`. It is one GEMM between two `permutedims`. +- The **tensor product** multiplies the slices of two tensors pairwise, + `_mul(x, y)[:, :, k] == x[:, :, k] * y[:, :, k]`, with the `tensor_tensor_mul` kernel. + +A `SymmetricMatrix`, `SkewSymMatrix`, `StrictlyLowerTriangular` or `StrictlyUpperTriangular` weight +multiplies a vector or a matrix through its own `*`. Before a GEMM with a tensor, and as the right +argument of a left product, it is made dense once per call, on its own backend, and its gradient +reaches the storage through `ProjectTo`. Element types are promoted as `*` promotes them. + +# Examples + +A left product: + +```jldoctest +using GeometricMachineLearning: _mul + +B = [1 1 1; 1 1 1; 1 1 1;;; 2 2 2; 2 2 2; 2 2 2] +A = [3 0 0; 0 2 0; 0 0 1] + +_mul(A, B) + +# output + +3×3×2 Array{Int64, 3}: +[:, :, 1] = + 3 3 3 + 2 2 2 + 1 1 1 + +[:, :, 2] = + 6 6 6 + 4 4 4 + 2 2 2 +``` + +A right product: + +```jldoctest +using GeometricMachineLearning: _mul + +A = [1 1 1; 1 1 1; 1 1 1;;; 2 2 2; 2 2 2; 2 2 2] +B = [3 0 0; 0 2 0; 0 0 1] + +_mul(A, B) + +# output + +3×3×2 Array{Int64, 3}: +[:, :, 1] = + 3 2 1 + 3 2 1 + 3 2 1 + +[:, :, 2] = + 6 4 2 + 6 4 2 + 6 4 2 +``` +""" +_mul(W::AbstractMatrix, x::AbstractVecOrMat) = W * _dense(x) + +function _mul(W::AbstractMatrix, x::AbstractArray) + reshape(_dense(W) * _matrix(x), size(W, 1), Base.tail(size(x))...) +end + +function _mul(x::AbstractArray{<:Number, 3}, W::AbstractMatrix) + xW = reshape(permutedims(x, (1, 3, 2)), :, size(x, 2)) * _dense(W) + permutedims(reshape(xW, size(x, 1), size(x, 3), size(W, 2)), (1, 3, 2)) +end + +_mul(x::AbstractArray{<:Number, 3}, y::AbstractArray{<:Number, 3}) = tensor_tensor_mul(x, y) + +# A structured matrix as a dense matrix on its own backend; `one` allocates the identity there. +function _dense(W::Union{ + SymmetricMatrix, SkewSymMatrix, StrictlyLowerTriangular, StrictlyUpperTriangular}) + W * one(W) +end +_dense(x::AbstractArray) = x + +# `x` as a dense matrix: its trailing axes folded into the second, or, for a structured matrix, +# `_dense(x)`. A matrix is not reshaped, which would hide an `Adjoint` behind a `ReshapedArray`. +_matrix(x::AbstractMatrix) = _dense(x) +_matrix(x::AbstractArray) = reshape(x, size(x, 1), :) + +# Without this rule a `Matrix` parameter `P` multiplied as `_mul(P', x)` gets an `Adjoint` +# gradient: the generic rule of `*` gives `W` a dense cotangent `C`, and the rule of `adjoint` hands +# `C'` to `P`. This rule gives `W` the cotangent `C'` instead, so that `P` receives `C`. +function ChainRulesCore.rrule(::typeof(_mul), W::Adjoint{<:Number, <:AbstractMatrix}, x::AbstractArray) + y = _mul(W, x) + project_x = ProjectTo(x) + function _mul_adjoint_pullback(ȳ) + ȳ = _seed(eltype(y), ȳ) + ȳ isa AbstractZero && return NoTangent(), ZeroTangent(), ZeroTangent() + ȳ_matrix = _matrix(ȳ) + x_matrix = _matrix(x) + W̄ = convert(AbstractMatrix{eltype(W)}, x_matrix * ȳ_matrix')' + x̄ = project_x(reshape(W' * ȳ_matrix, size(x))) + NoTangent(), W̄, x̄ + end + y, _mul_adjoint_pullback +end diff --git a/src/kernels/tensor_mat_mul.jl b/src/kernels/tensor_mat_mul.jl deleted file mode 100644 index 688e42ca7..000000000 --- a/src/kernels/tensor_mat_mul.jl +++ /dev/null @@ -1,122 +0,0 @@ -# `C[:, :, k] = A[:, :, k] * B`: every slice of a tensor times a matrix, one thread per output entry. -@kernel function tensor_mat_mul_kernel!(C, A, B) - i, j, k = @index(Global, NTuple) - - # creating a temporary sum variable for matrix multiplication - tmp_sum = zero(eltype(C)) - for l in 1:size(A)[2] - tmp_sum += A[i, l, k] * B[l, j] - end - - C[i, j, k] = tmp_sum -end - -@doc raw""" - tensor_mat_mul!(C, A, B) - -Multiply the matrix `B` onto the tensor `A` from the right and store the result in `C`. - -The function [`tensor_mat_mul`](@ref) calls `tensor_mat_mul!` internally. -""" -function tensor_mat_mul!(C::AbstractArray{<:Number, 3}, A::AbstractArray{<:Number, 3}, B::AbstractMatrix) - @assert size(A)[2] == size(B)[1] - - backend = networkbackend(A) - kernel! = tensor_mat_mul_kernel!(backend) - kernel!(C, A, B, ndrange = size(C)) -end - -@doc raw""" - tensor_mat_mul(A::AbstractArray{<:Number, 3}, B::AbstractMatrix) - -Multipliy the matrix `B` onto the tensor `A` from the right. - -Internally this calls the inplace version [`tensor_mat_mul!`](@ref). - -# Examples - -```jldoctest -using GeometricMachineLearning: tensor_mat_mul - -A = [1 1 1; 1 1 1; 1 1 1;;; 2 2 2; 2 2 2; 2 2 2] -B = [3 0 0; 0 2 0; 0 0 1] - -tensor_mat_mul(A, B) - -# output - -3×3×2 Array{Int64, 3}: -[:, :, 1] = - 3 2 1 - 3 2 1 - 3 2 1 - -[:, :, 2] = - 6 4 2 - 6 4 2 - 6 4 2 -``` -""" -function tensor_mat_mul(A::AbstractArray{<:Number, 3}, B::AbstractMatrix) - @assert eltype(A) == eltype(B) - T = eltype(A) - sizeA = size(A) - sizeB = size(B) - @assert sizeA[2] == sizeB[1] - tensor_shape = (sizeA[1], sizeB[2], sizeA[3]) - backend = networkbackend(A) - C = KernelAbstractions.allocate(backend, T, tensor_shape...) - tensor_mat_mul!(C, A, B) - C -end - -########################### SymmetricMatrix (right multiplication) - -@kernel function symmetric_mat_right_mul_kernel!( - C::AbstractArray{T, 3}, B::AbstractArray{T, 3}, - S::AbstractVector{T}, n::Int) where {T} - i, j, l = @index(Global, NTuple) - tmp_sum = zero(T) - - for k in j:n - tmp_sum += B[i, k, l] * S[(k - 1) * k ÷ 2 + j] - end - - for k in 1:(j - 1) - tmp_sum += B[i, k, l] * S[(j - 1) * j ÷ 2 + k] - end - - C[i, j, l] = tmp_sum -end - -function symmetric_mat_right_mul!(C::AbstractArray{T, 3}, B::AbstractArray{T, 3}, - S::AbstractVector{T}, n::Int) where {T} - backend = networkbackend(C) - - symmetric_mat_right_mul_k! = symmetric_mat_right_mul_kernel!(backend) - symmetric_mat_right_mul_k!(C, B, S, n, ndrange = size(C)) - - nothing -end - -function symmetric_mat_right_mul(B::AbstractArray{T, 3}, S::AbstractVector{T}, n::Int) where {T} - C = copy(B) - - symmetric_mat_right_mul!(C, B, S, n) - - C -end - -@doc raw""" - mat_tensor_mul!(C::AbstractArray{<:Number, 3}, B::AbstractArray{<:Number, 3}, A::SymmetricMatrix) - -Multiply the symmetric matrix `A` onto the tensor `B` from the right and store the result in `C`. - -This performs an efficient multiplication based on the special structure of the symmetric matrix `A`. -""" -function tensor_mat_mul!(C::AbstractArray{<:Number, 3}, B::AbstractArray{<:Number, 3}, A::SymmetricMatrix) - @assert eltype(C) == eltype(B) == eltype(A) - @assert A.n == size(C, 2) == size(B, 2) - - symmetric_mat_right_mul!(C, B, A.S, A.n) -end diff --git a/src/kernels/tensor_tensor_transpose_mul.jl b/src/kernels/tensor_tensor_transpose_mul.jl index 9172d881f..5c7996738 100644 --- a/src/kernels/tensor_tensor_transpose_mul.jl +++ b/src/kernels/tensor_tensor_transpose_mul.jl @@ -1,7 +1,3 @@ -""" -This implements the operation (A,B) -> A*B' for two tensors -""" - # `C[:, :, k] = A[:, :, k] * B[:, :, k]'`: slice by slice, one thread per output entry. @kernel function tensor_tensor_transpose_mul_kernel!(C, A, B) i, j, k = @index(Global, NTuple) diff --git a/src/kernels/tensor_transpose.jl b/src/kernels/tensor_transpose.jl deleted file mode 100644 index 0fca89c78..000000000 --- a/src/kernels/tensor_transpose.jl +++ /dev/null @@ -1,29 +0,0 @@ -""" -This computes the transpose of a tensor. -""" - -@kernel function tensor_transpose_kernel!(C, A) - i, j, k = @index(Global, NTuple) - - C[i, j, k] = A[j, i, k] -end - -function tensor_transpose!(C, A) - sizeA = size(A) - sizeC = size(C) - @assert sizeA[3] == sizeC[3] - @assert sizeA[1] == sizeC[2] - @assert sizeA[2] == sizeC[1] - - backend = networkbackend(A) - kernel! = tensor_transpose_kernel!(backend) - kernel!(C, A, ndrange = size(C)) -end - -function tensor_transpose(A::AbstractArray{T, 3}) where {T} - sizeA = size(A) - backend = networkbackend(A) - C = KernelAbstractions.zeros(backend, T, sizeA[2], sizeA[1], sizeA[3]) - tensor_transpose!(C, A) - C -end diff --git a/src/kernels/tensor_transpose_mat_mul.jl b/src/kernels/tensor_transpose_mat_mul.jl deleted file mode 100644 index d6cda274a..000000000 --- a/src/kernels/tensor_transpose_mat_mul.jl +++ /dev/null @@ -1,30 +0,0 @@ -""" -This implements the operation (A,B) -> A'*B for a tensor and a matrix -""" - -@kernel function tensor_transpose_mat_mul_kernel!(C, A, B) - i, j, k = @index(Global, NTuple) - - # creating a temporary sum variable for matrix multiplication - tmp_sum = zero(eltype(C)) - for l in 1:size(A)[1] - tmp_sum += A[l, i, k] * B[l, j] - end - - C[i, j, k] = tmp_sum -end - -function tensor_transpose_mat_mul!(C, A, B) - @assert axes(A, 1) == axes(B, 1) - - backend = networkbackend(A) - kernel! = tensor_transpose_mat_mul_kernel!(backend) - kernel!(C, A, B, ndrange = size(C)) -end - -function tensor_transpose_mat_mul(A::AbstractArray{T, 3}, B::AbstractMatrix{T}) where {T} - backend = networkbackend(A) - C = KernelAbstractions.allocate(backend, T, size(A)[2], size(B)[2], size(A)[3]) - tensor_transpose_mat_mul!(C, A, B) - C -end diff --git a/src/kernels/tensor_transpose_tensor_mul.jl b/src/kernels/tensor_transpose_tensor_mul.jl index 849626a05..516001548 100644 --- a/src/kernels/tensor_transpose_tensor_mul.jl +++ b/src/kernels/tensor_transpose_tensor_mul.jl @@ -1,7 +1,3 @@ -""" -This implements the operation (A,B) -> A'*B for two tensors -""" - # `C[:, :, k] = A[:, :, k]' * B[:, :, k]`: slice by slice, one thread per output entry. @kernel function tensor_transpose_tensor_mul_kernel!(C, A, B) i, j, k = @index(Global, NTuple) diff --git a/src/kernels/vec_tensor_mul.jl b/src/kernels/vec_tensor_mul.jl deleted file mode 100644 index 3bcf10db0..000000000 --- a/src/kernels/vec_tensor_mul.jl +++ /dev/null @@ -1,13 +0,0 @@ -@kernel function vec_tensor_mul_kernel!( - b::AbstractArray{T, 3}, a::AbstractVector{T}, x::AbstractArray{T, 3}) where {T} - i, j, k = @index(Global, NTuple) - b[i, j, k] = x[i, j, k]*a[i] -end - -function vec_tensor_mul(a::AbstractVector{T}, x::AbstractArray{T, 3}) where {T} - b = similar(x) - backend = networkbackend(x) - vec_tensor_mul! = vec_tensor_mul_kernel!(backend) - vec_tensor_mul!(b, a, x, ndrange = size(x)) - b -end diff --git a/src/layers/classification.jl b/src/layers/classification.jl index a02a8e183..9741ce880 100644 --- a/src/layers/classification.jl +++ b/src/layers/classification.jl @@ -81,7 +81,7 @@ end function (d::ClassificationLayer{M, N, true})(output::AbstractArray{T, 3}, ps::NamedTuple) where { M, N, T} - d.activation(sum(mat_tensor_mul(ps.weight, output), dims = 2)/size(output, 2)) + d.activation(sum(_mul(ps.weight, output), dims = 2)/size(output, 2)) end function (d::ClassificationLayer{M, N, true})(output::AbstractArray{T, 2}, ps::NamedTuple) where { @@ -91,7 +91,7 @@ end function (d::ClassificationLayer{M, N, false})(output::AbstractArray{T, 3}, ps::NamedTuple) where { M, N, T} - d.activation(assign_output_estimate(mat_tensor_mul(ps.weight, output), 1)) + d.activation(assign_output_estimate(_mul(ps.weight, output), 1)) end function (d::ClassificationLayer{M, N, false})(output::AbstractArray{T, 2}, ps::NamedTuple) where { diff --git a/src/layers/linear_symplectic_attention.jl b/src/layers/linear_symplectic_attention.jl index 4fddd4faf..9340c273c 100644 --- a/src/layers/linear_symplectic_attention.jl +++ b/src/layers/linear_symplectic_attention.jl @@ -80,16 +80,13 @@ end # Implement multiplication with symmetric matrix from the right! function (::LinearSymplecticAttentionQ)(z::NamedTuple{(:q, :p), Tuple{AT, AT}}, ps::NamedTuple) where {AT} - (q = z.q + _custom_mul(z.p, ps.A), p = z.p) + (q = z.q + _mul(z.p, ps.A), p = z.p) end function (::LinearSymplecticAttentionP)(z::NamedTuple{(:q, :p), Tuple{AT, AT}}, ps::NamedTuple) where {AT} - (q = z.q, p = z.p + _custom_mul(z.q, ps.A)) + (q = z.q, p = z.p + _mul(z.q, ps.A)) end function (d::LinearSymplecticAttention)(z::AbstractArray, ps::NamedTuple) apply_layer_to_nt_and_return_array(z, d, ps) end - -_custom_mul(z::AbstractArray{T, 3}, A::AbstractMatrix{T}) where {T} = tensor_mat_mul(z, A) -_custom_mul(z::AbstractMatrix{T}, A::AbstractMatrix{T}) where {T} = z * A diff --git a/src/layers/manifold_layer.jl b/src/layers/manifold_layer.jl index b0fe1b09d..7cb83d41d 100644 --- a/src/layers/manifold_layer.jl +++ b/src/layers/manifold_layer.jl @@ -1,10 +1,6 @@ # This defines a manifold layer that only has one matrix-valued manifold $A$ associated with it does $x\mapsto{}Ax$. abstract type ManifoldLayer{M, N} <: AbstractExplicitLayer{M, N} end -function (d::ManifoldLayer{M, N})(x::AbstractVecOrMat, ps::NamedTuple) where {M, N} - N > M ? ps.weight*x : ps.weight'*x -end - -function (d::ManifoldLayer{M, N})(x::AbstractArray{T, 3}, ps::NamedTuple) where {M, N, T} - N > M ? mat_tensor_mul(ps.weight, x) : mat_tensor_mul(ps.weight', x) +function (d::ManifoldLayer{M, N})(x::AbstractArray, ps::NamedTuple) where {M, N} + N > M ? _mul(ps.weight, x) : _mul(ps.weight', x) end diff --git a/src/layers/multi_head_attention.jl b/src/layers/multi_head_attention.jl index 8b8eebc5a..d846b20db 100644 --- a/src/layers/multi_head_attention.jl +++ b/src/layers/multi_head_attention.jl @@ -114,8 +114,8 @@ function compute_output_of_mha(d::MultiHeadAttention{M, M}, x::AbstractMatrix{T} head_outputs = map(1:(d.n_heads)) do i key = Symbol("head_", i) - ps.PV[key]' * x * - d.activation((ps.PQ[key]' * x)' * (ps.PK[key]' * x) / T(sqrt(dim))) + _mul(ps.PV[key]', x) * + d.activation(_mul(ps.PQ[key]', x)' * _mul(ps.PK[key]', x) / T(sqrt(dim))) end # One `vcat` over all heads, not one per head. The head outputs are matrices, so `reduce` takes @@ -134,9 +134,9 @@ function compute_output_of_mha( # the result of a single head attention block, one per head head_outputs = map(1:(d.n_heads)) do i key = Symbol("head_", i) - Q_tensor = mat_tensor_mul(ps.PQ[key]', x) - K_tensor = mat_tensor_mul(ps.PK[key]', x) - V_tensor = mat_tensor_mul(ps.PV[key]', x) + Q_tensor = _mul(ps.PQ[key]', x) + K_tensor = _mul(ps.PK[key]', x) + V_tensor = _mul(ps.PV[key]', x) QK_tensor = tensor_transpose_tensor_mul(Q_tensor, K_tensor) tensor_tensor_mul(V_tensor, d.activation(QK_tensor/T(sqrt(dim)))) @@ -161,19 +161,6 @@ function (d::MultiHeadAttention{M, M, Stiefel, false})(x::AbstractArray, ps::Nam compute_output_of_mha(d, x, ps) end -function mat_tensor_mul(Y::AT, - x::AbstractArray{T, 3}) where {T <: Number, - BT <: AbstractArray{T}, - ST <: StiefelManifold{T, BT}, - AT <: Adjoint{T, ST}} - mat_tensor_mul(Y.parent.A', x) -end - -# Multiply `Y` with every matrix stored in `x`, parallelised over the third axis. -function mat_tensor_mul(Y::StiefelManifold, x::AbstractArray{<:Number, 3}) - mat_tensor_mul(Y.A, x) -end - function (d::MultiHeadAttention)(z::QPT, ps::NamedTuple) N2 = size(z.q, 1) output = d(vcat(z.q, z.p), ps) diff --git a/src/layers/psd_like_layer.jl b/src/layers/psd_like_layer.jl index fa17812a3..e38de984b 100644 --- a/src/layers/psd_like_layer.jl +++ b/src/layers/psd_like_layer.jl @@ -38,8 +38,8 @@ end function (::PSDLayer{M, N})(qp::NamedTuple{(:q, :p), Tuple{AT1, AT2}}, ps::NamedTuple) where {M, N, AT1 <: AbstractArray, AT2 <: AbstractArray} - N > M ? (q = custom_mat_mul(ps.weight, qp.q), p = custom_mat_mul(ps.weight, qp.p)) : - (q = custom_mat_mul(ps.weight', qp.q), p = custom_mat_mul(ps.weight', qp.p)) + N > M ? (q = _mul(ps.weight, qp.q), p = _mul(ps.weight, qp.p)) : + (q = _mul(ps.weight', qp.q), p = _mul(ps.weight', qp.p)) end function (l::PSDLayer{M, N})(x::AbstractArray, ps::NamedTuple) where {M, N} diff --git a/src/layers/resnet.jl b/src/layers/resnet.jl index 6a6b109b7..322b27029 100644 --- a/src/layers/resnet.jl +++ b/src/layers/resnet.jl @@ -44,20 +44,12 @@ function parameterlength(::ResNetLayer{M, M, use_bias}) where {M, use_bias} use_bias ? M * (M + 1) : M * M end -function (d::ResNetLayer{M, M, true})(x::AbstractVecOrMat, ps::NamedTuple) where {M} - x + d.activation.(ps.weight * x .+ ps.bias) +function (d::ResNetLayer{M, M, true})(x::AbstractArray, ps::NamedTuple) where {M} + x + d.activation.(_mul(ps.weight, x) .+ ps.bias) end -function (d::ResNetLayer{M, M, false})(x::AbstractVecOrMat, ps::NamedTuple) where {M} - x + d.activation.(ps.weight * x) -end - -function (d::ResNetLayer{M, M, false})(x::AbstractArray{T, 3}, ps::NamedTuple) where {M, T} - x + d.activation.(mat_tensor_mul(ps.weight, x)) -end - -function (d::ResNetLayer{M, M, true})(x::AbstractArray{T, 3}, ps::NamedTuple) where {M, T} - x + d.activation.(mat_tensor_mul(ps.weight, x) .+ ps.bias) +function (d::ResNetLayer{M, M, false})(x::AbstractArray, ps::NamedTuple) where {M} + x + d.activation.(_mul(ps.weight, x)) end function (d::Union{Dense{M, N}, ResNetLayer{M, N}})(z::QPT, ps::NamedTuple) where {M, N} diff --git a/src/layers/symplectic_attention.jl b/src/layers/symplectic_attention.jl index a3ba6d7c9..ab933a8d0 100644 --- a/src/layers/symplectic_attention.jl +++ b/src/layers/symplectic_attention.jl @@ -120,11 +120,11 @@ end function (d::SymplecticAttentionQ{M, M, :arbitrary})( z::NamedTuple{(:q, :p), Tuple{AT, AT}}, ps::NamedTuple) where {AT, M} - PAP = _custom_mul(_custom_mul(_custom_transpose(z.p), ps.A), z.p) + PAP = _mul(_mul(_custom_transpose(z.p), ps.A), z.p) σPAP = d.activation(PAP) ( - q = z.q + (_custom_mul(_custom_mul(ps.A, z.p), _custom_transpose(σPAP)) + - _custom_mul(_custom_mul(ps.A', z.p), σPAP)), + q = z.q + (_mul(_mul(ps.A, z.p), _custom_transpose(σPAP)) + + _mul(_mul(ps.A', z.p), σPAP)), p = z.p) end @@ -136,37 +136,32 @@ end function (d::SymplecticAttentionQ{M, M, :symmetric})( z::NamedTuple{(:q, :p), Tuple{AT, AT}}, ps::NamedTuple) where {AT, M} A = ps.A - PAP = _custom_mul(_custom_mul(_custom_transpose(z.p), A), z.p) + PAP = _mul(_mul(_custom_transpose(z.p), A), z.p) σPAP = d.activation(PAP) - (q = z.q + _custom_mul(_custom_mul(A, z.p), 2 * σPAP), p = z.p) + (q = z.q + _mul(_mul(A, z.p), 2 * σPAP), p = z.p) end function (d::SymplecticAttentionP{M, M, :arbitrary})( z::NamedTuple{(:q, :p), Tuple{AT, AT}}, ps::NamedTuple) where {AT, M} - QAQ = _custom_mul(_custom_mul(_custom_transpose(z.q), ps.A), z.q) + QAQ = _mul(_mul(_custom_transpose(z.q), ps.A), z.q) σQAQ = d.activation(QAQ) (q = z.q, - p = z.p + (_custom_mul(_custom_mul(ps.A, z.q), _custom_transpose(σQAQ)) + - _custom_mul(_custom_mul(ps.A', z.q), σQAQ))) + p = z.p + (_mul(_mul(ps.A, z.q), _custom_transpose(σQAQ)) + + _mul(_mul(ps.A', z.q), σQAQ))) end # One access to `ps.A`, for the reason given at the `SymplecticAttentionQ` method above. function (d::SymplecticAttentionP{M, M, :symmetric})( z::NamedTuple{(:q, :p), Tuple{AT, AT}}, ps::NamedTuple) where {AT, M} A = ps.A - QAQ = _custom_mul(_custom_mul(_custom_transpose(z.q), A), z.q) + QAQ = _mul(_mul(_custom_transpose(z.q), A), z.q) σQAQ = d.activation(QAQ) - (q = z.q, p = z.p + _custom_mul(_custom_mul(A, z.q), 2 * σQAQ)) + (q = z.q, p = z.p + _mul(_mul(A, z.q), 2 * σQAQ)) end function (d::SymplecticAttention)(z::AbstractArray, ps::NamedTuple) apply_layer_to_nt_and_return_array(z, d, ps) end -_custom_mul(A::AbstractMatrix{T}, z::AbstractArray{T, 3}) where {T} = mat_tensor_mul(A, z) -function _custom_mul(Z1::AbstractArray{T, 3}, Z2::AbstractArray{T, 3}) where {T} - tensor_tensor_mul(Z1, Z2) -end - _custom_transpose(Z::AbstractMatrix) = Z' -_custom_transpose(z::AbstractArray{T, 3}) where {T} = tensor_transpose(z) +_custom_transpose(z::AbstractArray{<:Number, 3}) = permutedims(z, (2, 1, 3)) diff --git a/src/layers/sympnets.jl b/src/layers/sympnets.jl index ee7d6570f..a95ff29c1 100644 --- a/src/layers/sympnets.jl +++ b/src/layers/sympnets.jl @@ -215,54 +215,39 @@ function parameterlength(::LinearLayer{M, M}) where {M} (M ÷ 2) * (M ÷ 2 + 1) ÷ 2 end -# Multiplies a matrix with a vector, a matrix or a tensor. -custom_mat_mul(weight::AbstractMatrix, x::AbstractVecOrMat) = weight*x -function custom_mat_mul(weight::AbstractMatrix, x::AbstractArray{T, 3}) where {T} - mat_tensor_mul(weight, x) -end - -function custom_vec_mul(scale::AbstractVector, x::AbstractVecOrMat) - scale .* x -end -function custom_vec_mul(scale::AbstractVector{T}, x::AbstractArray{T, 3}) where {T} - vec_tensor_mul(scale, x) -end - @inline function (d::ActivationLayerQ{M, M})(x::NamedTuple, ps) where {M} size(x.q, 1) == M÷2 || error("Dimension mismatch.") - return (q = x.q + custom_vec_mul(ps.scale, d.activation.(x.p)), p = x.p) + return (q = x.q + ps.scale .* d.activation.(x.p), p = x.p) end @inline function (d::ActivationLayerP{M, M})(x::NamedTuple, ps) where {M} size(x.q, 1) == M÷2 || error("Dimension mismatch.") - return (q = x.q, p = x.p + custom_vec_mul(ps.scale, d.activation.(x.q))) + return (q = x.q, p = x.p + ps.scale .* d.activation.(x.q)) end @inline function (d::GradientLayerQ{M, M})(x::NamedTuple, ps) where {M} size(x.q, 1) == M÷2 || error("Dimension mismatch.") ( - q = x.q + custom_mat_mul(ps.weight', - (custom_vec_mul(ps.scale, d.activation.(custom_mat_mul(ps.weight, x.p) .+ - ps.bias)))), + q = x.q + _mul(ps.weight', + ps.scale .* d.activation.(_mul(ps.weight, x.p) .+ ps.bias)), p = x.p) end @inline function (d::GradientLayerP{M, M})(x::NamedTuple, ps) where {M} size(x.q, 1) == M÷2 || error("Dimension mismatch.") (q = x.q, - p = x.p + custom_mat_mul(ps.weight', - (custom_vec_mul(ps.scale, d.activation.(custom_mat_mul(ps.weight, x.q) .+ - ps.bias))))) + p = x.p + _mul(ps.weight', + ps.scale .* d.activation.(_mul(ps.weight, x.q) .+ ps.bias))) end @inline function (d::LinearLayerQ{M, M})(x::NamedTuple, ps) where {M} size(x.q, 1) == M÷2 || error("Dimension mismatch.") - (q = x.q + custom_mat_mul(ps.weight, x.p), p = x.p) + (q = x.q + _mul(ps.weight, x.p), p = x.p) end @inline function (d::LinearLayerP{M, M})(x::NamedTuple, ps) where {M} size(x.q, 1) == M÷2 || error("Dimension mismatch.") - (q = x.q, p = x.p + custom_mat_mul(ps.weight, x.q)) + (q = x.q, p = x.p + _mul(ps.weight, x.q)) end # For the wrappers whose input to a SympNet layer is an `AbstractArray` rather than the `NamedTuple` diff --git a/src/layers/volume_preserving_attention.jl b/src/layers/volume_preserving_attention.jl index 7a2592e62..2b7c4901f 100644 --- a/src/layers/volume_preserving_attention.jl +++ b/src/layers/volume_preserving_attention.jl @@ -95,7 +95,7 @@ function (d::VolumePreservingAttention{M, M, :skew_sym})( T, 3}, ps::NamedTuple) where {T, M} tensor_tensor_mul(x, orthonormal_activation_cayley(d, - tensor_transpose_tensor_mul(x, mat_tensor_mul(ps.A, x)) + tensor_transpose_tensor_mul(x, _mul(ps.A, x)) ) ) end @@ -105,7 +105,7 @@ function (d::VolumePreservingAttention{M, M, :arbitrary})( 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 - tensor_transpose(x_interim)) + orthonormal_activation_cayley(d, x_interim - permutedims(x_interim, (2, 1, 3))) ) end diff --git a/src/layers/volume_preserving_feedforward.jl b/src/layers/volume_preserving_feedforward.jl index ef81bfc5e..355f02c06 100644 --- a/src/layers/volume_preserving_feedforward.jl +++ b/src/layers/volume_preserving_feedforward.jl @@ -176,29 +176,10 @@ function initialparameters( (weight = StrictlyUpperTriangular(S, M),) end -function (d::VolumePreservingFeedForwardLayer{M, M, :bias})( - x::AbstractArray{ - T, 3}, ps) where {T, M} - x + d.activation.(mat_tensor_mul(ps.weight, x) .+ ps.bias) +function (d::VolumePreservingFeedForwardLayer{M, M, :bias})(x::AbstractArray, ps) where {M} + x + d.activation.(_mul(ps.weight, x) .+ ps.bias) end -function (d::VolumePreservingFeedForwardLayer{M, M, :no_bias})( - x::AbstractArray{ - T, 3}, ps) where {T, M} - x + d.activation.(mat_tensor_mul(ps.weight, x)) -end - -function (d::VolumePreservingFeedForwardLayer{M, M, :bias})(x::AbstractMatrix{T}, ps) where { - T, M} - x + d.activation.(ps.weight * x .+ ps.bias) -end - -function (d::VolumePreservingFeedForwardLayer{M, M, :bias})(x::AbstractVector{T}, ps) where { - T, M} - x + d.activation.(ps.weight * x + ps.bias) -end - -function (d::VolumePreservingFeedForwardLayer{M, M, :no_bias})(x::AbstractVecOrMat{T}, ps) where { - T, M} - x + d.activation.(ps.weight * x) +function (d::VolumePreservingFeedForwardLayer{M, M, :no_bias})(x::AbstractArray, ps) where {M} + x + d.activation.(_mul(ps.weight, x)) end diff --git a/test/Project.toml b/test/Project.toml index 3359d574b..95f4c0279 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -1,15 +1,18 @@ [deps] AbstractNeuralNetworks = "60874f82-5ada-4c70-bd1c-fa6be7711c8a" Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" +ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a" Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4" ExplicitImports = "7d51a73a-1435-4ff3-83d9-f097790105c7" FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000" +GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527" GeometricIntegrators = "dcce2d33-59f6-5b8d-9047-0defad88ae06" GeometricMachineLearning = "194d25b2-d3f5-49f0-af24-c124f4aa80cc" GeometricOptimizers = "fc236c15-5557-4942-aa65-b650f329279e" GeometricProblems = "18cb22b4-ad41-5c80-9c5f-710df63fbdc9" HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f" +JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb" JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b" KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" @@ -35,6 +38,7 @@ FiniteDifferences = "0.12" GeometricIntegrators = "0.18.2" GeometricProblems = "0.8, 0.9" HDF5 = "0.16, 0.17" +JLArrays = "0.3" Metal = "1.10" Printf = "1" SafeTestsets = "0.1" diff --git a/test/arrays/triangular.jl b/test/arrays/triangular.jl index abb754f66..4b64c778a 100644 --- a/test/arrays/triangular.jl +++ b/test/arrays/triangular.jl @@ -1,5 +1,5 @@ using GeometricMachineLearning -using GeometricMachineLearning: mat_tensor_mul +using GeometricMachineLearning: _mul using Zygote: pullback using Test import Random @@ -9,15 +9,15 @@ Random.seed!(1234) # What the triangular types *are* — their storage layout, their arithmetic, their multiplication # against a dense matrix — is tested in GeometricOptimizers, which defines them # (`test/special_matrices/triangular.jl` there). What is left here is GML's: batching them over the -# third axis of a tensor with `mat_tensor_mul`, and the pullback of that kernel. +# third axis of a tensor with `_mul`, and the pullback of that product. function triangular_tensor_multiplication_test(T = Float64, n::Int = 5) Aₗ = rand(StrictlyLowerTriangular{T}, n) Aᵤ = rand(StrictlyUpperTriangular{T}, n) B = rand(T, n, n, n) - AₗB = mat_tensor_mul(Aₗ, B) - AᵤB = mat_tensor_mul(Aᵤ, B) + AₗB = _mul(Aₗ, B) + AᵤB = _mul(Aᵤ, B) for i in 1:n @test AₗB[:, :, i] ≈ Aₗ * B[:, :, i] @test AᵤB[:, :, i] ≈ Aᵤ * B[:, :, i] @@ -31,8 +31,8 @@ function triangular_tensor_multiplication_pullback_test(T = Float64, n::Int = 5) B = rand(T, n, n, n) C_diff = rand(T, n, n, n) - total_pb_lower = pullback(mat_tensor_mul, Aₗ, B)[2](C_diff) - total_pb_upper = pullback(mat_tensor_mul, Aᵤ, B)[2](C_diff) + total_pb_lower = pullback(_mul, Aₗ, B)[2](C_diff) + total_pb_upper = pullback(_mul, Aᵤ, B)[2](C_diff) # The batched pullback has to agree slice by slice with the pullback of the single-slice # product. These were bare expressions and not `@test`s before, so the loop asserted nothing. diff --git a/test/devices/jlarray.jl b/test/devices/jlarray.jl new file mode 100644 index 000000000..ac08faf71 --- /dev/null +++ b/test/devices/jlarray.jl @@ -0,0 +1,61 @@ +# 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. +# Each result must stay on `JLBackend()` and match the host. + +using GeometricMachineLearning +using GeometricMachineLearning: _mul, tensor_tensor_mul, assign_q_and_p +using ChainRulesCore: rrule +using GPUArraysCore: allowscalar +using JLArrays: JLArray, JLBackend +import KernelAbstractions +using Zygote: gradient +using Test +import Random + +Random.seed!(123) +allowscalar(false) + +ondevice(x) = KernelAbstractions.get_backend(x) isa JLBackend + +@testset "tensor_tensor_mul and its pullback on a JLArray" begin + for T in (Float32, Float64) + a, b, c̄ = rand(T, 3, 4, 5), rand(T, 4, 2, 5), rand(T, 3, 2, 5) + c, pullback = rrule(tensor_tensor_mul, JLArray(a), JLArray(b)) + _, ā, b̄ = pullback(JLArray(c̄)) + _, ā_host, b̄_host = rrule(tensor_tensor_mul, a, b)[2](c̄) + @test all(ondevice, (c, ā, b̄)) + @test isapprox(Array(c), tensor_tensor_mul(a, b); rtol = sqrt(eps(T))) + @test isapprox(Array(ā), ā_host; rtol = sqrt(eps(T))) + @test isapprox(Array(b̄), b̄_host; rtol = sqrt(eps(T))) + end +end + +@testset "_mul with a structured weight on a JLArray: forward and storage gradient" begin + for T in (Float32, Float64) + x = rand(T, 4, 6, 3) + for TT in (SymmetricMatrix, SkewSymMatrix, StrictlyLowerTriangular, StrictlyUpperTriangular) + W = rand(TT{T}, 4) + W_device = TT(JLArray(W.S), W.n) + ps = NetworkParameters((L1 = (W = W,),)) + ps_device = NetworkParameters((L1 = (W = W_device,),)) + y = _mul(W_device, JLArray(x)) + g = gradient(p -> sum(abs2, _mul(p.L1.W, JLArray(x))), ps_device)[1].L1.W + g_host = gradient(p -> sum(abs2, _mul(p.L1.W, x)), ps)[1].L1.W + @test ondevice(y) + @test ondevice(g.S) + @test isapprox(Array(y), _mul(W, x); rtol = sqrt(eps(T))) + @test isapprox(Array(g.S), g_host.S; rtol = sqrt(eps(T))) + end + end +end + +@testset "assign_q_and_p on a JLArray and on a view of one" begin + for T in (Float32, Float64) + x = rand(T, 6, 3) + for v in (JLArray(x), view(JLArray(x), :, 1:2)) + qp = assign_q_and_p(v, 3) + @test ondevice(qp.q) && ondevice(qp.p) + @test Array(qp.p) == Array(v)[4:6, :] + end + end +end diff --git a/test/devices/metal.jl b/test/devices/metal.jl index 832258340..a12234279 100644 --- a/test/devices/metal.jl +++ b/test/devices/metal.jl @@ -13,7 +13,7 @@ # that reads an array one element at a time raises instead of answering slowly. using GeometricMachineLearning -using GeometricMachineLearning: tensor_mat_mul! +using GeometricMachineLearning: _mul, tensor_tensor_mul import KernelAbstractions using Metal using Test @@ -45,12 +45,20 @@ if Metal.functional() @test Array(adjoint_product) ≈ J * B' end - @testset "tensor_mat_mul! on Metal matches the host" begin - a = rand(T, 5, 12, 100) - b = rand(T, 12, 10) - c = KernelAbstractions.zeros(backend, T, 5, 10, 100) - tensor_mat_mul!(c, MtlArray(a), MtlArray(b)) - @test all(i -> Array(c)[:, :, i] ≈ a[:, :, i] * b, 1:100) + @testset "_mul and tensor_tensor_mul on Metal match the host" begin + x = rand(T, 12, 10, 100) + W = rand(T, 5, 12) + S = rand(SymmetricMatrix{T}, 12) + V = rand(T, 10, 4) + y = rand(T, 10, 7, 100) + products = ((_mul(MtlArray(W), MtlArray(x)), _mul(W, x)), + (_mul(SymmetricMatrix(MtlArray(S.S), S.n), MtlArray(x)), _mul(S, x)), + (_mul(MtlArray(x), MtlArray(V)), _mul(x, V)), + (tensor_tensor_mul(MtlArray(x), MtlArray(y)), tensor_tensor_mul(x, y))) + for (device, host) in products + @test KernelAbstractions.get_backend(device) == backend + @test Array(device) ≈ host + end end else @test_skip Metal.functional() diff --git a/test/integration/mixed_precision_training.jl b/test/integration/mixed_precision_training.jl new file mode 100644 index 000000000..838a24b5f --- /dev/null +++ b/test/integration/mixed_precision_training.jl @@ -0,0 +1,20 @@ +using GeometricMachineLearning +using Test +import Random + +Random.seed!(123) + +# A network with a Stiefel leaf trains on data of the other precision, as a network of `Dense` +# layers does: the products promote, and every leaf keeps the network's element type. +@testset "a Stiefel network trains on data of the other precision" begin + for (T, S) in ((Float64, Float32), (Float32, Float64)) + x = rand(S, 6, 12) + nn = NeuralNetwork(Chain(StiefelLayer(6, 3), Dense(3, 2, tanh)), T) + before = deepcopy(nn.params) + train!(nn, TrainingOptimizer(nn.params), DataLoader(x, x[1:2, :]), Batch(4), 2; + loss = FeedForwardLoss()) + leaves = (values(nn.params.L1)..., values(nn.params.L2)...) + @test all(leaf -> eltype(leaf) === T, leaves) + @test nn.params.L1.weight != before.L1.weight + end +end diff --git a/test/integration/product_gradients.jl b/test/integration/product_gradients.jl new file mode 100644 index 000000000..5f3b54d51 --- /dev/null +++ b/test/integration/product_gradients.jl @@ -0,0 +1,105 @@ +using GeometricMachineLearning +using GeometricMachineLearning: _mul, params +using FiniteDifferences: central_fdm, grad +using Zygote: gradient +using Test +import Random + +Random.seed!(123) + +# The finite-difference gradient of `loss(W)` with respect to the storage `W.S` of a structured +# weight, in `Float64` on the `Float64` copy of the data, which is what the gradient of a +# `NetworkParameters` holds in its structured leaf. +function storage_fd(loss, W) + grad(central_fdm(5, 1), S -> loss(typeof(W).name.wrapper(S, W.n)), Float64.(W.S))[1] +end + +# `Float64` at `1e-8`, above the `1e-12` the products reach; `Float32` at its own accuracy. +tolerance(T) = T === Float64 ? 1e-8 : sqrt(eps(Float32)) + +# The network's single structured leaf replaced by `W`, for the finite-difference loss. +with_leaf(ps, W) = NetworkParameters((L1 = NamedTuple{keys(ps.L1)}((W,)),)) + +@testset "a structured weight in _mul: the storage gradient equals finite differences" begin + for T in (Float32, Float64) + x = rand(T, 4, 5, 3) + z = rand(T, 3, 4, 5) + for TT in (SymmetricMatrix, SkewSymMatrix, StrictlyLowerTriangular, StrictlyUpperTriangular) + ps = NetworkParameters((L1 = (W = rand(TT{T}, 4),),)) + g = gradient(p -> sum(abs2, sin.(_mul(p.L1.W, x))), ps)[1].L1.W + fd = storage_fd(W -> sum(abs2, sin.(_mul(W, Float64.(x)))), ps.L1.W) + @test g isa TT{T} + @test isapprox(g.S, fd; rtol = tolerance(T)) + end + for TT in (SymmetricMatrix, SkewSymMatrix, StrictlyLowerTriangular, StrictlyUpperTriangular) + ps = NetworkParameters((L1 = (W = rand(TT{T}, 4),),)) + g = gradient(p -> sum(abs2, sin.(_mul(z, p.L1.W))), ps)[1].L1.W + fd = storage_fd(W -> sum(abs2, sin.(_mul(Float64.(z), W))), ps.L1.W) + @test g isa TT{T} + @test isapprox(g.S, fd; rtol = tolerance(T)) + end + end +end + +@testset "LinearLayerQ: the same gradient on a matrix and on the same data as a tensor" begin + for T in (Float32, Float64) + nn = NeuralNetwork(Chain(LinearLayerQ(4)), CPU(), T) + ps = params(nn) + x = rand(T, 4, 7) + g_matrix = gradient(p -> sum(abs2, nn.model(x, p)), ps)[1].L1.weight + g_tensor = gradient(p -> sum(abs2, nn.model(reshape(x, 4, 1, 7), p)), ps)[1].L1.weight + fd = storage_fd(W -> sum(abs2, nn.model(Float64.(x), with_leaf(ps, W))), ps.L1.weight) + @test g_matrix isa SymmetricMatrix{T} + @test g_tensor isa SymmetricMatrix{T} + @test isapprox(g_tensor.S, g_matrix.S; rtol = tolerance(T)) + @test isapprox(g_tensor.S, fd; rtol = tolerance(T)) + end +end + +@testset "the symplectic attention layers on a tensor: the storage gradient equals finite differences" begin + for T in (Float32, Float64), + (layer, x) in ((LinearSymplecticAttentionQ(4, 5), rand(T, 4, 5, 3)), + (SymplecticAttentionQ(4; symmetric = true), rand(T, 4, 3, 2))) + + nn = NeuralNetwork(Chain(layer), CPU(), T) + ps = params(nn) + g = only(values(gradient(p -> sum(abs2, nn.model(x, p)), ps)[1].L1)) + W = only(values(ps.L1)) + fd = storage_fd(W -> sum(abs2, nn.model(Float64.(x), with_leaf(ps, W))), W) + @test g isa SymmetricMatrix{T} + @test isapprox(g.S, fd; rtol = tolerance(T)) + end +end + +# A Stiefel weight multiplied as `Y'` gets a dense `Matrix` gradient leaf, on matrix input as on +# tensor input, and the two agree on the same data. The same holds for every head of a +# `MultiHeadAttention`, with Stiefel weights and with dense ones. +@testset "a Stiefel weight multiplied as its adjoint: the gradient leaf is a Matrix" begin + for T in (Float32, Float64) + nn = NeuralNetwork(Chain(StiefelLayer(6, 3)), CPU(), T) + ps = params(nn) + x = rand(T, 6, 5) + g_matrix = gradient(p -> sum(abs2, nn.model(x, p)), ps)[1].L1.weight + g_tensor = gradient(p -> sum(abs2, nn.model(reshape(x, 6, 1, 5), p)), ps)[1].L1.weight + @test g_matrix isa Matrix{T} + @test g_tensor isa Matrix{T} + @test isapprox(g_matrix, g_tensor; rtol = sqrt(eps(T))) + @test isapprox(g_matrix, 2 * x * (Matrix(ps.L1.weight)' * x)'; rtol = sqrt(eps(T))) + + for Stiefel in (true, false) + mha = NeuralNetwork(Chain(MultiHeadAttention(4, 2; Stiefel = Stiefel)), CPU(), T) + ps = params(mha) + x = rand(T, 4, 3) + g_matrix = gradient(p -> sum(abs2, mha.model(x, p)), ps)[1].L1 + g_tensor = gradient(p -> sum(abs2, mha.model(reshape(x, 4, 3, 1), p)), ps)[1].L1 + for g in (g_matrix, g_tensor) + @test all(head -> head isa Matrix{T}, ( + values(g.PQ)..., values(g.PK)..., values(g.PV)...)) + end + for key in (:PQ, :PK, :PV), head in keys(g_matrix[key]) + + @test isapprox(g_matrix[key][head], g_tensor[key][head]; rtol = sqrt(eps(T))) + end + end + end +end diff --git a/test/kernels/assign_q_and_p.jl b/test/kernels/assign_q_and_p.jl new file mode 100644 index 000000000..1417fb44c --- /dev/null +++ b/test/kernels/assign_q_and_p.jl @@ -0,0 +1,36 @@ +using GeometricMachineLearning: assign_q_and_p +using Zygote: gradient +using Test +import Random + +Random.seed!(123) + +rows(x, r) = x[r, ntuple(_ -> :, ndims(x) - 1)...] + +# A loss that reads one half of the split sees the gradient of that half and zeros on the other, in +# the shape of the input. +@testset "assign_q_and_p: the gradient of one half" begin + N = 3 + for T in (Float32, Float64), x in (rand(T, 2N), rand(T, 2N, 4), rand(T, 2N, 4, 5)) + + g = gradient(z -> sum(assign_q_and_p(z, N).p), x)[1] + @test g isa typeof(x) + @test g == vcat(zeros(T, size(rows(x, 1:N))), ones(T, size(rows(x, 1:N)))) + g = gradient(z -> sum(assign_q_and_p(z, N).q), x)[1] + @test g == vcat(ones(T, size(rows(x, 1:N))), zeros(T, size(rows(x, 1:N)))) + end +end + +@testset "assign_q_and_p: the halves are the first N rows and the N after them" begin + N = 2 + for T in (Float32, Float64), + x in (rand(T, 2N + 1), rand(T, 2N + 1, 3), rand(T, 2N + 1, 3, 4)) + + for n in (N, Int32(N)) + qp = assign_q_and_p(x, n) + @test qp.q == rows(x, 1:N) + @test qp.p == rows(x, (N + 1):(2N)) + @test qp.p isa typeof(x) + end + end +end diff --git a/test/kernels/inverses/test_cayley_transforms.jl b/test/kernels/inverses/test_cayley_transforms.jl index 6369c148d..9c4b085a1 100644 --- a/test/kernels/inverses/test_cayley_transforms.jl +++ b/test/kernels/inverses/test_cayley_transforms.jl @@ -1,5 +1,4 @@ -using GeometricMachineLearning: tensor_cayley4, tensor_cayley3, cpu_tensor_cayley, - tensor_transpose +using GeometricMachineLearning: tensor_cayley4, tensor_cayley3, cpu_tensor_cayley using Test import Random @@ -20,8 +19,8 @@ function test_cayley_transforms(third_dim::Int = 10) A₄ = rand(4, 4, third_dim) A₃ = rand(3, 3, third_dim) # asymmetrize - A₄ = A₄ - tensor_transpose(A₄) - A₃ = A₃ - tensor_transpose(A₃) + 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₄)) diff --git a/test/kernels/kernel_pullbacks.jl b/test/kernels/kernel_pullbacks.jl index 24997c813..a276ba318 100644 --- a/test/kernels/kernel_pullbacks.jl +++ b/test/kernels/kernel_pullbacks.jl @@ -1,48 +1,36 @@ -using GeometricMachineLearning: lo_mat_mul, up_mat_mul, skew_mat_mul, symmetric_mat_mul, - symmetric_mat_right_mul -using GeometricMachineLearning: tensor_mat_mul, mat_tensor_mul, tensor_tensor_mul, - tensor_transpose_tensor_mul, assign_q_and_p, - tensor_transpose, assign_output_estimate, vec_tensor_mul, - tensor_mat_skew_sym_assign, tensor_transpose_mat_mul +using GeometricMachineLearning: tensor_tensor_mul, tensor_transpose_tensor_mul, + assign_output_estimate, + tensor_mat_skew_sym_assign, cpu_inverse, tensor_inverse2, + tensor_inverse3, tensor_inverse4, tensor_inverse5, + init_output, + _mul +using ChainRulesCore: rrule, AbstractZero, NoTangent, ZeroTangent, @thunk using ChainRulesTestUtils -using Printf +using FiniteDifferences: central_fdm, j′vp +using LinearAlgebra: I +using Test import Random 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) + function main(first_dim, second_dim, third_dim, third_tensor_dim) - # the first two are tests for the splitting x -> (q, p) (vector and matrix) - test_rrule(assign_q_and_p, rand(first_dim*2), first_dim) - test_rrule(assign_q_and_p, rand(first_dim*2, second_dim), first_dim) - test_rrule(assign_q_and_p, rand(first_dim*2, second_dim, third_dim), first_dim) - test_rrule(tensor_mat_mul, rand(first_dim, second_dim, third_tensor_dim), rand(second_dim, third_dim)) - test_rrule(tensor_transpose_mat_mul, rand(second_dim, first_dim, third_tensor_dim), - rand(second_dim, third_dim)) - # test_rrule(tensor_transpose_mat_mul, rand(second_dim, first_dim, third_tensor_dim), rand(second_dim, third_dim)) - test_rrule(mat_tensor_mul, rand(first_dim, second_dim), rand(second_dim, third_dim, third_tensor_dim)) test_rrule(tensor_tensor_mul, rand(first_dim, second_dim, third_tensor_dim), rand(second_dim, third_dim, third_tensor_dim)) test_rrule(tensor_transpose_tensor_mul, rand(second_dim, first_dim, third_tensor_dim), rand(second_dim, third_dim, third_tensor_dim)) - test_rrule(tensor_transpose, rand(first_dim, second_dim, third_tensor_dim)) test_rrule(assign_output_estimate, rand(first_dim, second_dim, third_tensor_dim), 1) - test_rrule(vec_tensor_mul, rand(first_dim), rand(first_dim, second_dim, third_tensor_dim)) - - # mat_tensor_mul routines for special arrays - test_rrule(lo_mat_mul, rand(first_dim * (first_dim - 1) ÷ 2), - rand(first_dim, first_dim, third_dim), first_dim) - test_rrule(up_mat_mul, rand(first_dim * (first_dim - 1) ÷ 2), - rand(first_dim, first_dim, third_dim), first_dim) - test_rrule(skew_mat_mul, rand(first_dim * (first_dim - 1) ÷ 2), - rand(first_dim, first_dim, third_dim), first_dim) - test_rrule(symmetric_mat_mul, rand(first_dim * (first_dim + 1) ÷ 2), - rand(first_dim, second_dim, third_dim), first_dim) - test_rrule(symmetric_mat_right_mul, rand(second_dim, first_dim, third_dim), - rand(first_dim * (first_dim + 1) ÷ 2), first_dim) - test_rrule(tensor_mat_skew_sym_assign, rand(first_dim, second_dim, third_tensor_dim), + test_rrule( + tensor_mat_skew_sym_assign, rand(first_dim, second_dim, third_tensor_dim + 1), 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)) end const dim_range = 10 @@ -57,3 +45,82 @@ for _ in 1:num_tests main(first_dim, second_dim, third_dim, third_tensor_dim) 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),))) +end + +@testset "a zero seed gives a zero cotangent for every argument" begin + for T in (Float32, Float64), (f, args) in rules(T), + seed in (ZeroTangent(), NoTangent(), + @thunk(ZeroTangent())) + c̄ = rrule(f, args...)[2](seed) + @test c̄ isa Tuple + @test length(c̄) == 1 + length(args) + @test all(c -> c isa AbstractZero, c̄) + end +end + +@testset "a seed of the other precision gives cotangents of the primal's element type" begin + for (T, S) in ((Float32, Float64), (Float64, Float32)), (f, args) in rules(T) + + y, pullback = rrule(f, args...) + 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 + end + end +end + +# The seed is converted before the pullback computes, so a `Float64` seed gives exactly the +# cotangents of its `Float32` rounding. The last row cancels: computed in `Float64`, the cotangent +# of `W` keeps a residue of `eps(Float32) / 4` that the rounded seed does not have. +@testset "a seed of the other precision is converted at entry" begin + for (f, args) in rules(Float32) + y, pullback = rrule(f, args...) + ȳ = Float64.(rand(Float32, size(y))) .+ eps(Float32) / 4 + @test pullback(ȳ) == pullback(Float32.(ȳ)) + end + W = ones(Float32, 2, 1) + x = Float32[1 -1; 0 0] + ȳ = [1 + Float64(eps(Float32)) / 4 1.0] + pullback = rrule(_mul, W', x)[2] + @test ȳ isa Matrix{Float64} + @test pullback(ȳ) == pullback(Float32.(ȳ)) +end + +# The `Float64` rules are checked by `test_rrule` above; the `Float32` ones against finite +# differences in `Float64` on the same data, at the accuracy of `Float32`. +@testset "Float32 cotangents against Float64 finite differences" begin + for (f, args) in rules(Float32) + y, pullback = rrule(f, args...) + ȳ = rand(Float32, size(y)) + c̄ = Base.tail(pullback(ȳ)) + args64 = map(a -> Float64.(a), args) + fd = j′vp(central_fdm(5, 1), f, Float64.(ȳ), args64...) + for (c, reference) in zip(c̄, fd) + if c isa AbstractZero + @test maximum(abs, reference) < sqrt(eps(Float32)) + else + @test eltype(c) === Float32 + @test isapprox(c, reference; rtol = sqrt(eps(Float32))) + end + end + end +end diff --git a/test/kernels/mul.jl b/test/kernels/mul.jl new file mode 100644 index 000000000..910c83036 --- /dev/null +++ b/test/kernels/mul.jl @@ -0,0 +1,76 @@ +using GeometricMachineLearning +using GeometricMachineLearning: _mul +using Zygote: gradient +using Test +import Random + +Random.seed!(123) + +const STRUCTURED = ( + SymmetricMatrix, SkewSymMatrix, StrictlyLowerTriangular, StrictlyUpperTriangular) + +# `_mul(W, x)[:, :, k] == W * x[:, :, k]`, slice by slice, with the structured weight multiplied as a +# dense matrix. +function slices(f, x::AbstractArray{<:Number, 3}) + cat((f(x[:, :, k]) for k in axes(x, 3))...; dims = 3) +end + +@testset "_mul: left, right and tensor products against slice-by-slice references" begin + for T in (Float32, Float64) + x = rand(T, 4, 5, 3) + for W in (rand(T, 6, 4), rand(T, 4, 6)', (rand(TT{T}, 4) for TT in STRUCTURED)...) + Wd = Matrix(W) + y = @inferred _mul(W, x) + @test y isa Array{T, 3} + @test y ≈ slices(xk -> Wd * xk, x) + @test _mul(W, x[:, :, 1]) ≈ Wd * x[:, :, 1] + @test _mul(W, x[:, 1, 1]) ≈ Wd * x[:, 1, 1] + @test _mul(W, x[:, :, 1]) isa Matrix{T} + @test _mul(W, x[:, 1, 1]) isa Vector{T} + end + z = rand(T, 3, 4, 5) + for W in (rand(T, 4, 6), (rand(TT{T}, 4) for TT in STRUCTURED)...) + y = @inferred _mul(z, W) + @test y isa Array{T, 3} + @test y ≈ slices(zk -> zk * Matrix(W), z) + end + y = rand(T, 5, 2, 3) + @test @inferred(_mul(x, y)) isa Array{T, 3} + @test _mul(x, y) ≈ cat((x[:, :, k] * y[:, :, k] for k in 1:3)...; dims = 3) + end +end + +# A matrix argument on the right of a left product is multiplied as it stands, not through a +# reshape, which would turn a structured matrix into a `ReshapedArray`. +@testset "_mul: a matrix times a structured matrix" begin + for T in (Float32, Float64), TT in STRUCTURED + + z = rand(T, 3, 4) + A = rand(TT{T}, 4) + @test _mul(z, A) isa Matrix{T} + @test _mul(z, A) ≈ z * Matrix(A) + end +end + +# GEMM promotes, so a weight and data of two precisions multiply; no element type is checked. +@testset "_mul: mixed element types promote" begin + W = rand(Float32, 3, 4) + x = rand(Float64, 4, 2, 5) + @test _mul(W, x) isa Array{Float64, 3} + @test _mul(Float64.(W), Float32.(x)) isa Array{Float64, 3} + @test _mul(W', rand(Float64, 3, 2, 5)) isa Array{Float64, 3} +end + +# The cotangent of a dense `Matrix` parameter multiplied as `W'` is a `Matrix`, not an `Adjoint`, +# at every rank of the data. +@testset "_mul: the cotangent of W in _mul(W', x) is a Matrix" begin + for T in (Float32, Float64) + W = rand(T, 4, 3) + for x in (rand(T, 4), rand(T, 4, 5), rand(T, 4, 5, 2)) + g = gradient(w -> sum(abs2, _mul(w', x)), W)[1] + reference = gradient(w -> sum(abs2, Matrix(w') * reshape(x, 4, :)), W)[1] + @test g isa Matrix{T} + @test isapprox(g, reference; rtol = sqrt(eps(T))) + end + end +end diff --git a/test/kernels/tensor_cayley.jl b/test/kernels/tensor_cayley.jl index 9f9424e99..e30a55dab 100644 --- a/test/kernels/tensor_cayley.jl +++ b/test/kernels/tensor_cayley.jl @@ -1,6 +1,5 @@ using GeometricMachineLearning: tensor_cayley2, tensor_cayley3, tensor_cayley4, tensor_cayley5 -using GeometricMachineLearning: tensor_transpose using Test import Random @@ -15,25 +14,25 @@ end function test_tensor_cayley2(T::Type = Float32, third_dim::Int = 10) A = rand(T, 2, 2, third_dim) - B = tensor_cayley2(A - tensor_transpose(A)) + 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 - tensor_transpose(A)) + 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 - tensor_transpose(A)) + 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 - tensor_transpose(A)) + B = tensor_cayley5(A - permutedims(A, (2, 1, 3))) check_orthonormal_property(B) end diff --git a/test/kernels/tensor_mat_mul.jl b/test/kernels/tensor_mat_mul.jl deleted file mode 100644 index 2ab1891e6..000000000 --- a/test/kernels/tensor_mat_mul.jl +++ /dev/null @@ -1,41 +0,0 @@ -using GeometricMachineLearning -using GeometricMachineLearning: tensor_mat_mul!, tensor_mat_mul, allocate -import KernelAbstractions -using Random, Test - -Random.seed!(123) - -backend = CPU() - -const dim1 = 5 -const dim2 = 12 -const dim3 = 10 -const num_data = 100 - -a = rand!(allocate(backend, Float32, dim1, dim2, num_data)) -b = rand!(allocate(backend, Float32, dim2, dim3)) -c = KernelAbstractions.zeros(backend, Float32, dim1, dim3, num_data) -tensor_mat_mul!(c, a, b) - -c_manual = KernelAbstractions.zeros(backend, Float32, dim1, dim3, num_data) -for i in 1:num_data - @views c_manual[:, :, i] = a[:, :, i] * b -end - -@test isapprox(c, c_manual) - -function test_tensor_multiplication( - first_dim::Int = dim1, second_dim::Int = dim2, third_dim::Int = dim3; T = Float64) - A = rand(SymmetricMatrix{T}, second_dim) - B = rand(first_dim, second_dim, third_dim) - BA = tensor_mat_mul(B, A) - for l in 1:third_dim - @test (@view BA[:, :, l]) ≈ (@view B[:, :, l]) * A - end -end - -test_tensor_multiplication() - -A = [1 1 1; 1 1 1; 1 1 1;;; 2 2 2; 2 2 2; 2 2 2] -B = [3 0 0; 0 2 0; 0 0 1] -@test tensor_mat_mul(A, B) == [3 2 1; 3 2 1; 3 2 1;;; 6 4 2; 6 4 2; 6 4 2] diff --git a/test/kernels/test_skew_map.jl b/test/kernels/test_skew_map.jl index 7ef78c65e..a434ce46d 100644 --- a/test/kernels/test_skew_map.jl +++ b/test/kernels/test_skew_map.jl @@ -1,6 +1,6 @@ using GeometricMachineLearning -using GeometricMachineLearning: mat_tensor_mul, tensor_mat_skew_sym_assign, - tensor_transpose_tensor_mul, tensor_transpose +using GeometricMachineLearning: _mul, tensor_mat_skew_sym_assign, + tensor_transpose_tensor_mul using Test import Random @@ -20,7 +20,7 @@ end function test_first_option(n::Int = 10, seq_length::Int = 10, batch_size::Int = 20) A = rand(SkewSymMatrix, n) x = rand(n, seq_length, batch_size) - xAx = tensor_transpose_tensor_mul(x, mat_tensor_mul(A, x)) + xAx = tensor_transpose_tensor_mul(x, _mul(A, x)) isskew(xAx) end @@ -28,7 +28,7 @@ function test_second_option(n::Int = 10, seq_length::Int = 10, batch_size::Int = A = rand(n, n) x = rand(n, seq_length, batch_size) xAx = tensor_mat_skew_sym_assign(x, A) / √n - isskew(xAx - tensor_transpose(xAx)) + isskew(xAx - permutedims(xAx, (2, 1, 3))) end test_first_option() diff --git a/test/layers/double_multiplication_network_parameters_gradient.jl b/test/layers/double_multiplication_network_parameters_gradient.jl index 98c1c2cfe..b7d75f22c 100644 --- a/test/layers/double_multiplication_network_parameters_gradient.jl +++ b/test/layers/double_multiplication_network_parameters_gradient.jl @@ -1,6 +1,6 @@ using Test using GeometricMachineLearning -using GeometricMachineLearning: _custom_mul +using GeometricMachineLearning: _mul import Random include("../helpers/network_parameters_gradient_structure.jl") @@ -9,19 +9,19 @@ Random.seed!(1234) function single_multiplication(a::AT, _ps::Union{ NamedTuple, NetworkParameters}) where {AT <: AbstractArray} - _custom_mul(a, _ps.L1.A) + _mul(a, _ps.L1.A) end function double_multiplication(a::AT, _ps::Union{ NamedTuple, NetworkParameters}) where {AT <: AbstractArray} - _custom_mul(_ps.L1.A, _custom_mul(a, _ps.L1.A)) + _mul(_ps.L1.A, _mul(a, _ps.L1.A)) end S = rand(SymmetricMatrix, 4) ps = NetworkParameters((L1 = (A = S,),)) t = rand(4, 4) -@testset "_custom_mul: NetworkParameters gradient structure" begin +@testset "_mul: NetworkParameters gradient structure" begin @testset "single_multiplication (one wrapper access)" begin network_parameters_gradient_structure_test( _ps -> single_multiplication(t, _ps), ps) diff --git a/test/layers/symplectic_attention_network_parameters_gradient.jl b/test/layers/symplectic_attention_network_parameters_gradient.jl index 450e9b100..206aec5d9 100644 --- a/test/layers/symplectic_attention_network_parameters_gradient.jl +++ b/test/layers/symplectic_attention_network_parameters_gradient.jl @@ -1,6 +1,6 @@ using Test using GeometricMachineLearning -using GeometricMachineLearning: _custom_mul, _custom_transpose, params +using GeometricMachineLearning: _mul, _custom_transpose, params import Random include("../helpers/network_parameters_gradient_structure.jl") @@ -9,18 +9,18 @@ Random.seed!(1234) function symplectic_attention(z::NamedTuple{(:q, :p), Tuple{AT, AT}}, _ps::Union{NamedTuple, NetworkParameters}) where {AT <: AbstractArray} - expPAP = exp.(_custom_mul(_custom_mul(_custom_transpose(z.p), _ps.L1.A), z.p)) - (q = z.q + _custom_mul(_custom_mul(_ps.L1.A, z.p), 2 * expPAP) / sum(expPAP), p = z.p) + expPAP = exp.(_mul(_mul(_custom_transpose(z.p), _ps.L1.A), z.p)) + (q = z.q + _mul(_mul(_ps.L1.A, z.p), 2 * expPAP) / sum(expPAP), p = z.p) end function symplectic_attention_simplified(z::NamedTuple{(:q, :p), Tuple{AT, AT}}, _ps::Union{NamedTuple, NetworkParameters}) where {AT <: AbstractArray} - (q = z.p + _custom_mul(_custom_mul(z.p, _ps.L1.A), z.p), p = z.p) + (q = z.p + _mul(_mul(z.p, _ps.L1.A), z.p), p = z.p) end function symplectic_linear_map(z::NamedTuple{(:q, :p), Tuple{AT, AT}}, _ps::Union{NamedTuple, NetworkParameters}) where {AT <: AbstractArray} - (q = z.q + _custom_mul(_ps.L1.A, z.p), p = z.p) + (q = z.q + _mul(_ps.L1.A, z.p), p = z.p) end S = rand(SymmetricMatrix, 2) diff --git a/test/pullbacks/zygote_pullback.jl b/test/pullbacks/zygote_pullback.jl index 4ea9afd8c..598f9168c 100644 --- a/test/pullbacks/zygote_pullback.jl +++ b/test/pullbacks/zygote_pullback.jl @@ -1,7 +1,7 @@ using GeometricMachineLearning using GeometricMachineLearning: FeedForwardLoss, QPTOAT, LASympNet, StiefelLayer using AbstractNeuralNetworks: AbstractExplicitLayer, Chain, Dense, NeuralNetwork, params -using LinearAlgebra: Adjoint, Diagonal, diag +using LinearAlgebra: Diagonal, diag using NeuralNetworkParameters: NetworkParameters using Test import Random @@ -51,11 +51,12 @@ end @test eltype(leaf) === T @test size(leaf) == size(p) end - # The Stiefel leaf's tangent is an `Adjoint` of its storage rather than a `StiefelManifold`: - # the manifold layer's functor reads the weight through an adjoint, and Zygote keeps the - # cotangent's own type instead of projecting it back onto the leaf's. + # The Stiefel leaf's tangent is a dense `Matrix` of its storage rather than a + # `StiefelManifold`: the manifold layer reads the weight through an adjoint, whose product + # `_mul` hands the parent a dense cotangent, and Zygote keeps that type instead of + # projecting it back onto the leaf's. if network.model isa Chain && network.model.layers[1] isa StiefelLayer - @test g.L1.weight isa Adjoint{T, Matrix{T}} + @test g.L1.weight isa Matrix{T} @test size(g.L1.weight) == size(ps.L1.weight) end end diff --git a/test/quality/jet.jl b/test/quality/jet.jl index 33425d07b..98873f94b 100644 --- a/test/quality/jet.jl +++ b/test/quality/jet.jl @@ -40,43 +40,7 @@ const A3{T} = Array{T, 3} @testset "JET" begin if isdefined(JET, :JET_AVAILABLE) ? JET.JET_AVAILABLE : JET.JET_LOADABLE - # src/kernels/assign_q_and_p.jl: `test_rrule` in test/kernels/kernel_pullbacks.jl - @test isempty(JET.get_reports(JET.report_opt(GML.assign_q_and_p, - (Vector{Float64}, Int); target_modules = TARGET))) - @test isempty(JET.get_reports(JET.report_opt(GML.assign_q_and_p, - (Matrix{Float64}, Int); target_modules = TARGET))) - @test isempty(JET.get_reports(JET.report_opt(GML.assign_q_and_p, - (A3{Float64}, Int); target_modules = TARGET))) - @test isempty(kernel_body_reports(GML.assign_first_half!, 2, zeros(2), zeros(4))) - @test isempty(kernel_body_reports( - GML.assign_second_half!, 2, zeros(2), zeros(4), 2)) - @test isempty(kernel_body_reports(GML.assign_first_half!, (2, 3), zeros(2, 3), - zeros(4, 3))) - @test isempty(kernel_body_reports(GML.assign_second_half!, (2, 3), zeros(2, 3), - zeros(4, 3), 2)) - @test isempty(kernel_body_reports( - GML.assign_first_half!, (2, 3, 2), zeros(2, 3, 2), - zeros(4, 3, 2))) - @test isempty(kernel_body_reports(GML.assign_second_half!, (2, 3, 2), - zeros(2, 3, 2), zeros(4, 3, 2), 2)) - - # src/kernels/tensor_mat_mul.jl: `tensor_mat_mul!` in test/kernels/tensor_mat_mul.jl - @test isempty(JET.get_reports(JET.report_opt(GML.tensor_mat_mul!, - (A3{Float32}, A3{Float32}, Matrix{Float32}); target_modules = TARGET))) - @test isempty(kernel_body_reports(GML.tensor_mat_mul_kernel!, (3, 4, 2), - zeros(Float32, 3, 4, 2), zeros(Float32, 3, 5, 2), zeros(Float32, 5, 4))) - # `symmetric_mat_right_mul!` through `tensor_mat_mul!` with a `SymmetricMatrix`: - # `Float32` in test/architectures/linear_symplectic_transformer.jl, `Float64` in - # test/kernels/tensor_mat_mul.jl - for T in (Float32, Float64) - @test isempty(JET.get_reports(JET.report_opt(GML.symmetric_mat_right_mul!, - (A3{T}, A3{T}, Vector{T}, Int); target_modules = TARGET))) - @test isempty(kernel_body_reports( - GML.symmetric_mat_right_mul_kernel!, (2, 3, 2), - zeros(T, 2, 3, 2), zeros(T, 2, 3, 2), zeros(T, 6), 3)) - end - - # src/kernels/tensor_tensor_mul.jl and the three transposed products + # src/kernels/tensor_tensor_mul.jl and the two transposed products for T in (Float16, Float32, Float64) @test isempty(JET.get_reports(JET.report_opt(GML.tensor_tensor_mul!, (A3{T}, A3{T}, A3{T}); target_modules = TARGET))) @@ -93,83 +57,18 @@ const A3{T} = Array{T, 3} @test isempty(kernel_body_reports(GML.tensor_tensor_transpose_mul_kernel!, (2, 3, 2), zeros(T, 2, 3, 2), zeros(T, 2, 4, 2), zeros(T, 3, 4, 2))) end - @test isempty(JET.get_reports(JET.report_opt(GML.tensor_transpose_mat_mul!, - (A3{Float64}, A3{Float64}, Matrix{Float64}); target_modules = TARGET))) - @test isempty(kernel_body_reports(GML.tensor_transpose_mat_mul_kernel!, (2, 3, 2), - zeros(2, 3, 2), zeros(4, 2, 2), zeros(4, 3))) - # src/kernels/tensor_transpose.jl and src/kernels/exponentials/tensor_exponential.jl + # src/kernels/exponentials/tensor_exponential.jl for T in (Float16, Float32, Float64) - @test isempty(JET.get_reports(JET.report_opt(GML.tensor_transpose!, - (A3{T}, A3{T}); target_modules = TARGET))) - @test isempty(kernel_body_reports(GML.tensor_transpose_kernel!, (3, 2, 2), - zeros(T, 3, 2, 2), zeros(T, 2, 3, 2))) @test isempty(JET.get_reports(JET.report_opt(GML.assign_ones!, (A3{T},); target_modules = TARGET))) @test isempty(kernel_body_reports(GML.assign_ones_kernel!, (3, 2), zeros(T, 3, 3, 2))) end - # src/kernels/vec_tensor_mul.jl: `test_rrule` in test/kernels/kernel_pullbacks.jl - @test isempty(JET.get_reports(JET.report_opt(GML.vec_tensor_mul, - (Vector{Float64}, A3{Float64}); target_modules = TARGET))) - @test isempty(kernel_body_reports(GML.vec_tensor_mul_kernel!, (2, 3, 2), - zeros(2, 3, 2), zeros(2), zeros(2, 3, 2))) - - # src/kernels/mat_tensor_mul.jl - for T in (Float32, Float64) - @test isempty(JET.get_reports(JET.report_opt(GML.mat_tensor_mul!, - (A3{T}, Matrix{T}, A3{T}); target_modules = TARGET))) - @test isempty(kernel_body_reports(GML.mat_tensor_mul_kernel!, (2, 3, 2), - zeros(T, 2, 3, 2), zeros(T, 2, 4), zeros(T, 4, 3, 2))) - @test isempty(JET.get_reports(JET.report_opt(GML.symmetric_mat_mul!, - (A3{T}, Vector{T}, A3{T}, Int); target_modules = TARGET))) - @test isempty(kernel_body_reports(GML.symmetric_mat_mul_kernel!, (3, 2, 2), - zeros(T, 3, 2, 2), zeros(T, 6), zeros(T, 3, 2, 2), 3)) - end - @test isempty(JET.get_reports(JET.report_opt(GML.lo_mat_mul!, - (A3{Float64}, Vector{Float64}, A3{Float64}, Int); target_modules = TARGET))) - @test isempty(kernel_body_reports(GML.lo_mul_kernel!, (3, 2, 2), zeros(3, 2, 2), - zeros(3), zeros(3, 2, 2), 3)) - @test isempty(JET.get_reports(JET.report_opt(GML.up_mat_mul!, - (A3{Float64}, Vector{Float64}, A3{Float64}, Int); target_modules = TARGET))) - @test isempty(kernel_body_reports(GML.up_mul_kernel!, (3, 2, 2), zeros(3, 2, 2), - zeros(3), zeros(3, 2, 2), 3)) - for T in (Float16, Float32, Float64) - @test isempty(JET.get_reports(JET.report_opt(GML.skew_mat_mul!, - (A3{T}, Vector{T}, A3{T}, Int); target_modules = TARGET))) - @test isempty(kernel_body_reports(GML.skew_mat_mul_kernel!, (3, 2, 2), - zeros(T, 3, 2, 2), zeros(T, 3), zeros(T, 3, 2, 2), 3)) - end - - # src/kernels/kernel_ad_routines/: the pullbacks that launch kernels, at the `Float64` - # types of `test_rrule` in test/kernels/kernel_pullbacks.jl and of - # test/optimizers/structured_array_parameters.jl - S, Ss, B, dC = zeros(3), zeros(6), zeros(3, 3, 2), zeros(3, 3, 2) - for f in (GML.lo_mat_mul, GML.up_mat_mul, GML.skew_mat_mul) - @test isempty(JET.get_reports(JET.report_opt(rrule(f, S, B, 3)[2], - (A3{Float64},); target_modules = TARGET))) - end - @test isempty(JET.get_reports(JET.report_opt( - rrule(GML.symmetric_mat_mul, Ss, B, 3)[2], - (A3{Float64},); target_modules = TARGET))) - @test isempty(JET.get_reports(JET.report_opt( - rrule(GML.symmetric_mat_right_mul, B, Ss, 3)[2], (A3{Float64},); - target_modules = TARGET))) - @test isempty(kernel_body_reports(GML.lower_da_kernel!, size(dC), zero(dC), S, dC)) - @test isempty(kernel_body_reports(GML.lower_ds_kernel!, (3, 2), zeros(3, 2), B, dC)) - @test isempty(kernel_body_reports(GML.upper_da_kernel!, size(dC), zero(dC), S, dC)) - @test isempty(kernel_body_reports(GML.upper_ds_kernel!, (3, 2), zeros(3, 2), B, dC)) - @test isempty(kernel_body_reports(GML.symmetric_da_kernel!, size(dC), zero(dC), Ss, - dC)) - @test isempty(kernel_body_reports(GML.symmetric_ds_kernel!, (6, 2), zeros(6, 2), B, - dC)) - @test isempty(kernel_body_reports( - GML.symmetric_right_da_kernel!, size(dC), zero(dC), - Ss, dC)) - @test isempty(kernel_body_reports( - GML.symmetric_right_ds_kernel!, (6, 2), zeros(6, 2), - B, dC)) + # 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},); @@ -178,12 +77,6 @@ const A3{T} = Array{T, 3} zeros(4, 4, 2), 3, 4)) @test isempty(kernel_body_reports(GML.da_kernel!, (3, 3, 2), zeros(3, 3, 2), Z, M, zeros(4, 4, 2), 3, 4)) - for T in (Float32, Float64) - @test isempty(JET.get_reports(JET.report_opt(GML.tensor_scalar_product, - (A3{T}, A3{T}); target_modules = TARGET))) - @test isempty(kernel_body_reports(GML.tensor_scalar_product_kernel!, 3, - zeros(T, 3), zeros(T, 3, 4, 2), zeros(T, 3, 4, 2), 4, 2)) - end # 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},); diff --git a/test/runtests.jl b/test/runtests.jl index 66a2cceb2..60bcb4856 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -10,8 +10,12 @@ if "core" in GROUPS @safetestset "MatrixSoftmax" include("activations/softmax.jl") @safetestset "Symplectic Potential (array tests)" include("arrays/poisson_tensor.jl") @safetestset "Test triangular matrices" include("arrays/triangular.jl") - @safetestset "Custom tensor matrix multiplication" include("kernels/tensor_mat_mul.jl") + @safetestset "_mul: left, right and tensor products" include("kernels/mul.jl") + @safetestset "assign_q_and_p splits by slicing" include("kernels/assign_q_and_p.jl") @safetestset "Custom AD rules for kernels" include("kernels/kernel_pullbacks.jl") + @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 "Volume-Preserving Transformer (skew-symmetric tests)" include("kernels/test_skew_map.jl") @@ -59,7 +63,7 @@ if "core" in GROUPS @safetestset "Gradient optimizer tests" include("optimizers/gradient_optimizer.jl") @safetestset "Momentum optimizer tests" include("optimizers/momentum_optimizer.jl") @safetestset "Optimizers with structured (non-manifold) weights" include("optimizers/structured_array_parameters.jl") - @safetestset "_custom_mul: NetworkParameters gradient structure" include("layers/double_multiplication_network_parameters_gradient.jl") + @safetestset "_mul: NetworkParameters gradient structure" include("layers/double_multiplication_network_parameters_gradient.jl") @safetestset "Symplectic attention: NetworkParameters gradient structure" include("layers/symplectic_attention_network_parameters_gradient.jl") @safetestset "ZygotePullback returns a NetworkParameters" include("pullbacks/zygote_pullback.jl") @safetestset "SymbolicPullback of HNNLoss is the gradient of the batched loss" include("pullbacks/symbolic_hnn_pullback.jl")