diff --git a/CHANGELOG.md b/CHANGELOG.md index 714d8aa00..f2fc4308e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -55,6 +55,8 @@ breaking release). - **`LowerTriangular` and `UpperTriangular` are renamed `StrictlyLowerTriangular` and `StrictlyUpperTriangular`**, as in GeometricOptimizers 0.9, and are re-exported under the new names. A file saved with the old names no longer loads through the registry. The rigid-body tutorial registers the two old names for its committed files. +- **`ResNetLayer` encodes its bias as a `Symbol` type parameter, the two volume-preserving feedforward layers are one struct, and two fields and a keyword are renamed.** `ResNetLayer`'s third type parameter is `:bias` or `:no_bias`, as in the volume-preserving feedforward layers, where it was `true` or `false`; its keyword `use_bias::Bool` stays. `VolumePreservingFeedForwardLayer{M, N, bias, AT, Triangle}` is now the one struct of the two volume-preserving feedforward layers, with `Triangle` `:lower` or `:upper`, where it was their abstract supertype; `VolumePreservingLowerLayer{M, N, bias, AT}` and `VolumePreservingUpperLayer{M, N, bias, AT}` are its aliases and keep their constructors. Their `initialparameters` takes the bias initializer as the keyword `init_bias`, as `ResNetLayer` and the gradient layers do, where it took `init_bias!`. `GSympNet` stores its activation in the field `activation`, as `LASympNet` does, where it was `act`. `StandardTransformerIntegrator`'s field `upsacling_activation` is spelled `upscaling_activation`. `GradientLayer` stores `second_dim` as an `Int`, and the six SympNet layer constructors (`GradientLayerQ`/`P`, `ActivationLayerQ`/`P`, `LinearLayerQ`/`P`) convert the dimension to an `Int`, so `GradientLayerQ(Int32(4), Int32(8), tanh) === GradientLayerQ(4, 8, tanh)` and `ActivationLayerQ(Int32(4), tanh) === ActivationLayerQ(4, tanh)`. + - **The `[compat]` bounds move to the versions the training API needs**: GeometricOptimizers `"0.9"`, AbstractNeuralNetworks `"0.9.2"`, NeuralNetworkParameters `"0.5"`, SymbolicNeuralNetworks `"0.9"`, GeometricBase `"0.15"`, GeometricEquations `"0.21.5"` and GeometricSolutions `"0.6.6"`. HDF5 leaves `[compat]` and `[weakdeps]`. ### Fixed @@ -81,8 +83,16 @@ breaking release). - **A `DataLoader`'s two output fields are concrete, its minibatch index set is concrete and draws every sample, and a loader of one sample batches.** `output_dim` and `output_time_steps` were `Union{Int, Nothing}` whatever the loader held, which made a kernel launch read a value of no single type (issue #211); the copies of a solution and of an ensemble both go through one helper, `data_tensor`, which builds a 1-based tensor, and the batch indices are one host-built `2 × n` matrix copied to the backend once, with no `@views` (issue #212, which asked for `@views`). A `DataLoader` of one sample drew no minibatch (known issue K8) and a loader of matrices never drew its last sample (known issue K12); both now draw every pair, and the count of a trajectory is `cld(n · n_params, batch_size)`, so a seeded run draws the minibatches it drew before, except the ones the last sample joins. The `JET.report_opt` line of `test/quality/jet.jl` for a launcher that reads a `Union{Int, Nothing}` field is a `@test` rather than a `@test_broken`: the batch half of issue #325 is fixed here. This closes known issues B10, C17, K8 and K12, which leave `KNOWN_ISSUES.md`. B10's `DataLoader(::EnsembleSolution{T, T1, Vector{ST}})` for a `dataser` of `(t, q, q̇)` is reachable — a rigid-body `odeensemble` produces one, and the tutorial and three scripts call it — and `test/data_loader/data_loader.jl` tests it, so the method stays. +- **`LinearSymplecticTransformer` keeps the activation it is given.** Its constructor built `new{typeof(tanh)}`, so `LinearSymplecticTransformer(4, 3; activation = sin)` raised a `MethodError`; the gradient layers of its chain now carry `sin`. The test is `test/architectures/architecture_structure.jl`. + +- **`PositionalEncoding` allocates its encoding on the backend of its input.** The layer built a host `Matrix` and added it to its input, which a GPU kernel cannot read. It now fills `similar(x, eltype(x), dim, seq_length)` by one broadcast over the indices, so the layer runs on a `JLArray` with scalar indexing disallowed, and the matrix it adds is a `JLArray` there; `positional_encoding(T, dim, seq_length)` fills a host `Matrix` the same way, and the layer's output matches that of `4c4a8320` at 16 ULPs of the largest magnitude (`test/layers/layer_pins.jl`). The "CPU only" warnings of the `PositionalEncoding` docstring and of the `positional_encoding` keyword of `Transformer` are gone. This closes known issues B11 and K5, which leave `KNOWN_ISSUES.md`. The test is `test/layers/positional_encoding.jl`. + - **`assign_output_estimate` is a crop of the last `prediction_window` columns, and its gradient is Zygote's `getindex` adjoint.** Its kernel, `augment_zeros` and the `rrule` are gone: the adjoint is zero on the cropped columns, runs on a device with scalar indexing off, and gives a `Float32` cotangent for a `Float64` seed on `Float32` data, where `augment_zeros` returned the seed's type. The docstring's formula starts at column `T − pw + 1`, which is what the code does. +- **The decoder of an `AutoEncoder` stores the decoder's block count in `n_decoder_blocks`, where it stored the encoder's.** `decoder` passed `arch.n_encoder_blocks` to `UnknownDecoder` and `UnknownSymplecticDecoder`, so an `AutoEncoder` with two encoder blocks and three decoder blocks had the decoder architecture `UnknownDecoder(6, 2, 2)`; it is now `UnknownDecoder(6, 2, 3)`. The decoder's model and parameters already used the decoder's count, so no output changes, and the stored field changes only where the two counts differ. This affects every `AutoEncoder` that is not a `SymplecticAutoencoder`, `PSDArch` included. This closes known issue K24, which leaves `KNOWN_ISSUES.md`. The test is `test/architectures/architecture_pins.jl`, testset "encoder and decoder types of each autoencoder family". + +- **A `(q, p)` input whose `p` has the wrong size raises the layer's own `ErrorException("Dimension mismatch.")`.** The forward methods of `ActivationLayer`, `GradientLayer` and `LinearLayer` checked the size of `q` and not the size of `p`, so a `p` of the wrong size raised whichever error its first operation produced: a `DimensionMismatch` from a matrix product or a broadcast, or an `AssertionError` from the `SymmetricMatrix` product. Each method now checks `q` and `p` together. No valid input changes. This closes known issue K23, which leaves `KNOWN_ISSUES.md`. The test is `test/layers/layer_structure.jl`, testset "a SympNet layer raises on a half of the wrong size". + ### 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. @@ -143,6 +153,8 @@ breaking release). | `tensor_mat_skew_sym_assign` pullback (C15) | Float32 | — | 8365.5 → 8315.0 | | `tensor_mat_skew_sym_assign` pullback (C15) | Float64 | — | 8387.2 → 8354.8 | +- **A pair of layers or architectures that differed in one choice shares one method, and the outputs do not change.** One function, `_shift(Val(C), z, f)`, adds `f` of one half of a `(q, p)` input to the other half, so `src/layers/sympnets.jl` has three forward methods on a `(q, p)` input where it had six, `src/layers/symplectic_attention.jl` two where it had four, and `src/layers/linear_symplectic_attention.jl` one where it had two. Each `Bool` keyword of `VolumePreservingAttention`, the two volume-preserving feedforward layers, `SymplecticAttentionQ`/`SymplecticAttentionP` and `SymplecticAutoencoder` picks its type parameter in one expression. The two volume-preserving feedforward layers have one `initialparameters` where they had four. `MultiHeadAttention` has one `initialparameters` for both weight kinds, and it draws in the same order as before: head by head, and in each head the query, the key and the value projection, so no seeded value moves. Its forward pass still finds each head by its name `head_i`, but takes the names from a tuple that is built once for each number of heads, where it built `Symbol("head_", i)` for every head on every pass; at `4c4a8320`, 1000 passes of a two-head layer allocated 384 000 B more than the same products with constant keys (Julia 1.13.2). `LASympNet` has one `Chain` method where it had four; one `_gradient_layers` builds the alternating gradient layers of `GSympNet`, `LinearSymplecticTransformer` and `SymplecticTransformer`, and replaces the two methods of `_make_block_for_initialization`; the `DataLoader` constructors of `LASympNet`, `GSympNet` and `ClassificationTransformer` call the dimension constructors; and one `encoder` and one `decoder` serve every `AutoEncoder` that is not a `SymplecticAutoencoder`. Every type argument under `src/layers/` is spelled `::Type{T}`, and the aliases `LinearLayerQ` and `LinearLayerP` no longer name an unused third parameter, which Julia drops from an alias anyway. The outputs of 27 layer cases, 13 architecture cases and the encoders and decoders of two autoencoders, each in `Float32` and `Float64`, are compared at 16 ULPs of the largest magnitude of the pinned vector with the outputs that `4c4a8320` gives for the same seeds and inputs, in `test/layers/layer_pins.jl` and `test/architectures/architecture_pins.jl`; `test/architectures/forward_inference.jl` checks that the forward pass of every exported architecture with a chain infers, `StandardTransformerIntegrator(32)` with its default of 32 heads among them. This closes known issue C19, which leaves `KNOWN_ISSUES.md`. + - **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 diff --git a/KNOWN_ISSUES.md b/KNOWN_ISSUES.md index f6b73927f..e43f11ad6 100644 --- a/KNOWN_ISSUES.md +++ b/KNOWN_ISSUES.md @@ -2,28 +2,6 @@ What is known to be wrong in GeometricMachineLearning and is not fixed. -## B. Known defects - -### B11 · `PositionalEncoding` is CPU-only, and the failure is a compile error rather than a slowdown. - -- **location:** — -- **kind:** defect -- **found:** 2026-09-18 -- **evidence:** - `positional_encoding` builds its matrix with `Matrix{T}(undef, …)`, so the layer's - `x .+ P` adds a host array to whatever it is given. Every other layer allocates through the - backend — `KernelAbstractions.allocate`, or `similar(x, …)` in the forward pass. Broadcasting a - host `Matrix` against a device array does not fall back to the CPU; it fails to compile, because - the host array cannot be read from a kernel. - - **No test can catch this, because the suite has no GPU test.** That is what makes it worth an - entry rather than a docstring alone: a green matrix says nothing about it. Both docstrings — the - layer's and the `positional_encoding` keyword of [`Transformer`](@ref) — carry the warning. - - Closing it means giving the builder the backend, so that it allocates through - `KernelAbstractions.allocate` and fills with a kernel instead of a loop, and it cannot be verified - here without a GPU job to run it under. - ## C. Follow-ups and cleanups ### C6 · Three generated MNIST PDFs are in this branch's history @@ -89,20 +67,6 @@ What is known to be wrong in GeometricMachineLearning and is not fixed. deliberately, and the resolution is upstream declaring them `public`. Each check's reason is written beside it in the test file. -### C19 · `_make_block_for_initialization` has two methods - -- **location:** `src/architectures/symplectic_transformer.jl` -- **kind:** dead code -- **found:** 2026-09-19 -- **evidence:** - — `src/architectures/symplectic_transformer.jl` - and `src/architectures/linear_symplectic_transformer.jl` — that differ only in which fields they read - from their argument: one reads `arch.transformer_dim` and `arch.sympnet_activation`, the other reads - `arch.dim` and `arch.activation`. The bodies are otherwise identical. This stood as a `TODO` in the - source; a comment must stand on its own, so it is recorded here instead. Closing it means agreeing - on one set of field names across the two architectures, or passing the dimension and the activation - in rather than reading them off `arch`. - ### C20 · There is no Knet.jl example for the Hamiltonian neural network. - **location:** `TODO.md` @@ -261,6 +225,82 @@ What is known to be wrong in GeometricMachineLearning and is not fixed. `src/data_loader/data_loader.jl:205` states the one-type signature. Closing it means widening the three signatures to `NamedTupleOfArrays{T, (:q, :p)}`, and the docstring with them. +### K21 · `LagrangianNeuralNetwork` stores its activation in the field `act`, where every other architecture with one activation names it `activation`. + +- **location:** `src/architectures/lagrangian_neural_network.jl:5` +- **kind:** found late +- **found:** 2026-10-20, by the critics of part M5; the field is the same on `origin/main` at `4c4a8320` +- **evidence:** + ``` + $ grep -rn " act::\|\.act\b" src/architectures/ + src/architectures/lagrangian_neural_network.jl:5: act::AT + src/architectures/lagrangian_neural_network.jl:16: [Dense(nn.width, nn.width, nn.act) for _ in 1:nn.nhidden] + src/architectures/lagrangian_neural_network.jl:20: Dense(nn.dimin, nn.width, nn.act), + ``` + + The constructor's keyword is `activation`; only the field keeps the short name. The fix renames + the field, which changes a public field name. + +### K22 · A comment in `symplectic_attention.jl` names a test file under `test/parameters/`, which does not exist. + +- **location:** `src/layers/symplectic_attention.jl:131` +- **kind:** docs +- **found:** 2026-10-20, by the critics of part M5; the comment is the same on `origin/main` at `4c4a8320` +- **evidence:** + The comment above the symmetric forward method of `SymplecticAttention` names + `test/parameters/symplectic_attention_network_parameters_gradient.jl`. The file is + `test/layers/symplectic_attention_network_parameters_gradient.jl`, and `test/parameters/` does not + exist. The fix changes the path in the comment. + +### K25 · The docstrings of `LinearSymplecticTransformer` and `SymplecticTransformer` say that `n_sympnet = 1` gives one `GradientLayerQ` and one `GradientLayerP`. + +- **location:** `src/architectures/linear_symplectic_transformer.jl:22` +- **kind:** docs +- **found:** 2026-10-21, by the branch verifier of part M5; the docstrings and the chains are the same on `origin/main` at `4c4a8320` +- **evidence:** + Each transformer unit puts the same `n_sympnet` gradient layers after its `Q` attention layer and + after its `P` attention layer, and the alternation starts again in each block. With + `n_sympnet = 1` and the default `init_upper = true`, both gradient layers are `GradientLayerQ`: + + ```julia + [typeof(l) for l in Chain(LinearSymplecticTransformer(4, 3; n_sympnet = 1)).layers] + # LinearSymplecticAttentionQ, GradientLayerQ, LinearSymplecticAttentionP, GradientLayerQ + [typeof(l) for l in Chain(SymplecticTransformer(4; n_sympnet = 1)).layers] + # SymplecticAttentionQ, GradientLayerQ, SymplecticAttentionP, GradientLayerQ + ``` + + The same sentence is at `src/architectures/symplectic_transformer.jl:35`. The count + `2n_sympnet` holds for one unit, `L = 1`. The fix states the layers of one unit as the chain + builds them, or changes the chain, which changes the network. + +### K26 · The docstring of `SymplecticAutoencoder` says that `decoder_init_q = true` makes the first layer of each decoder block a ``p``-type layer, where it is a `GradientLayerQ`. + +- **location:** `src/architectures/symplectic_autoencoder.jl:36` +- **kind:** docs +- **found:** 2026-10-21, by the branch verifier of part M5; the docstring is the same on `origin/main` at `4c4a8320` +- **evidence:** + `decoder_init_q = true` selects `:DecoderInitQ`, whose `decoder_layers_from_iteration` passes + `isodd`, so layer 1 of each decoder block is a `GradientLayerQ`: + + ```julia + ae = NeuralNetwork(SymplecticAutoencoder(8, 4; decoder_init_q = true), CPU(), Float64) + first(decoder(ae).model.layers) isa GradientLayerQ # true + ``` + + The fix writes ``q`` for ``p`` in the docstring. + +### K27 · The docstring of `UnknownDecoder` names its third argument `n_encoder_blocks` and links to itself for an example. + +- **location:** `src/architectures/autoencoder.jl:102` +- **kind:** docs +- **found:** 2026-10-21, by the branch verifier of part M5; the docstring is the same on `origin/main` at `4c4a8320` +- **evidence:** + The signature line reads `UnknownDecoder(full_dim, reduced_dim, n_encoder_blocks)`, and the third + argument fills the field `n_decoder_blocks` (`src/architectures/autoencoder.jl:113`). The last + line says that an example "can be constructed analogously to [`UnknownDecoder`](@ref)", which + links the docstring to itself; the example is in the docstring of `UnknownEncoder`. The fix names + the argument `n_decoder_blocks` and links `UnknownEncoder`. + ## Upstream ### K2 · Revise prints EMFILE errors in the test log @@ -301,16 +341,6 @@ What is known to be wrong in GeometricMachineLearning and is not fixed. ## Documentation -### K5 · The `PositionalEncoding` docstring names `test/metal/`, which does not exist - -- **location:** `src/layers/positional_encoding.jl:75` -- **kind:** docs -- **found:** 2026-10-02 -- **evidence:** - The docstring says "the GPU tests in `test/metal/` do not cover this layer". The Metal tests - are in `test/devices/metal.jl`, and `origin/main` at 97bf0d8 has no `test/metal/` directory. - The fix is a change of that path under `src/`. - ### K18 · The tensor-sampling formula in the `DataLoader` manual is one short - **location:** `docs/src/data_loader/data_loader.md:143` diff --git a/src/architectures/autoencoder.jl b/src/architectures/autoencoder.jl index f48ce5e45..0801cdf36 100644 --- a/src/architectures/autoencoder.jl +++ b/src/architectures/autoencoder.jl @@ -184,15 +184,23 @@ function Chain(arch::AutoEncoder) Chain(encoder_model(arch).layers..., decoder_model(arch).layers...) end +# The architectures of the encoder and the decoder of an autoencoder that has none of its own: the +# symplectic ones for a `SymplecticCompression`. +_unknown_encoder(::AutoEncoder) = UnknownEncoder +_unknown_encoder(::SymplecticCompression) = UnknownSymplecticEncoder +_unknown_decoder(::AutoEncoder) = UnknownDecoder +_unknown_decoder(::SymplecticCompression) = UnknownSymplecticDecoder + """ encoder(nn::NeuralNetwork{<:AutoEncoder}) -Obtain the *encoder* from an [`AutoEncoder`](@ref) neural network. +Obtain the *encoder* from an [`AutoEncoder`](@ref) neural network. """ function encoder(nn::NeuralNetwork{<:AutoEncoder}) + arch = nn.architecture NeuralNetwork( - UnknownEncoder(nn.architecture.full_dim, nn.architecture.reduced_dim, nn.architecture.n_encoder_blocks), - encoder_model(nn.architecture), + _unknown_encoder(arch)(arch.full_dim, arch.reduced_dim, arch.n_encoder_blocks), + encoder_model(arch), encoder_parameters(nn), networkbackend(nn)) end @@ -223,9 +231,10 @@ end Obtain the *decoder* from an [`AutoEncoder`](@ref) neural network. """ function decoder(nn::NeuralNetwork{<:AutoEncoder}) + arch = nn.architecture NeuralNetwork( - UnknownDecoder(nn.architecture.full_dim, nn.architecture.reduced_dim, nn.architecture.n_encoder_blocks), - decoder_model(nn.architecture), + _unknown_decoder(arch)(arch.full_dim, arch.reduced_dim, arch.n_decoder_blocks), + decoder_model(arch), decoder_parameters(nn), networkbackend(nn)) end @@ -247,21 +256,3 @@ Internally this allocates a new nerual network of type [`UnknownDecoder`](@ref) function decoder(nn::NeuralNetwork) _decoder(nn, input_dimension(nn.model.layers[1]), output_dimension(nn.model.layers[end])) end - -function encoder(nn::NeuralNetwork{<:SymplecticCompression}) - NeuralNetwork( - UnknownSymplecticEncoder(nn.architecture.full_dim, nn.architecture.reduced_dim, - nn.architecture.n_encoder_blocks), - encoder_model(nn.architecture), - encoder_parameters(nn), - networkbackend(nn)) -end - -function decoder(nn::NeuralNetwork{<:SymplecticCompression}) - NeuralNetwork( - UnknownSymplecticDecoder(nn.architecture.full_dim, nn.architecture.reduced_dim, - nn.architecture.n_encoder_blocks), - decoder_model(nn.architecture), - decoder_parameters(nn), - networkbackend(nn)) -end diff --git a/src/architectures/linear_symplectic_transformer.jl b/src/architectures/linear_symplectic_transformer.jl index 8446dab71..180963d89 100644 --- a/src/architectures/linear_symplectic_transformer.jl +++ b/src/architectures/linear_symplectic_transformer.jl @@ -36,36 +36,19 @@ struct LinearSymplecticTransformer{AT} <: TransformerIntegrator where {AT} L::Int = lst_L_default, activation = lst_activation_default, init_upper::Bool = lst_init_upper_default) - new{typeof(tanh)}( + new{typeof(activation)}( dim, seq_length, n_sympnet, upscaling_dimension, L, activation, init_upper) end end -function _make_block_for_initialization( - layers::Tuple, arch::LinearSymplecticTransformer{AT}, - is_upper_criterion::Function) where {AT} - for i in 1:arch.n_sympnet - layers = (layers..., - if is_upper_criterion(i) - GradientLayerQ(arch.dim, arch.upscaling_dimension, arch.activation) - else - GradientLayerP(arch.dim, arch.upscaling_dimension, arch.activation) - end - ) - end - - layers -end - function Chain(arch::LinearSymplecticTransformer{AT}) where {AT} - # if `isodd(i)` and `arch.init_upper==true`, then do `GradientLayerQ`. - is_upper_criterion = arch.init_upper ? isodd : iseven + gradient_layers = _gradient_layers( + arch.dim, arch.upscaling_dimension, arch.activation, arch.n_sympnet, arch.init_upper) layers = () for _ in 1:arch.L - layers = (layers..., LinearSymplecticAttentionQ(arch.dim, arch.seq_length)) - layers = _make_block_for_initialization(layers, arch, is_upper_criterion) - layers = (layers..., LinearSymplecticAttentionP(arch.dim, arch.seq_length)) - layers = _make_block_for_initialization(layers, arch, is_upper_criterion) + layers = (layers..., LinearSymplecticAttentionQ(arch.dim, arch.seq_length), + gradient_layers..., LinearSymplecticAttentionP(arch.dim, arch.seq_length), + gradient_layers...) end Chain(layers...) diff --git a/src/architectures/standard_transformer_integrator.jl b/src/architectures/standard_transformer_integrator.jl index f5427d329..8e3658119 100644 --- a/src/architectures/standard_transformer_integrator.jl +++ b/src/architectures/standard_transformer_integrator.jl @@ -37,7 +37,7 @@ struct StandardTransformerIntegrator{AT1, AT2, AT3} <: TransformerIntegrator n_heads::Int n_blocks::Int L::Int - upsacling_activation::AT1 + upscaling_activation::AT1 resnet_activation::AT2 attention_activation::AT3 add_connection::Bool @@ -58,7 +58,7 @@ end function Chain(arch::StandardTransformerIntegrator) layers = arch.sys_dim == arch.transformer_dim ? () : - (Dense(arch.sys_dim, arch.transformer_dim, arch.upsacling_activation),) + (Dense(arch.sys_dim, arch.transformer_dim, arch.upscaling_activation),) for _ in 1:arch.L layers = (layers..., MultiHeadAttention( diff --git a/src/architectures/symplectic_autoencoder.jl b/src/architectures/symplectic_autoencoder.jl index a89243448..694713a91 100644 --- a/src/architectures/symplectic_autoencoder.jl +++ b/src/architectures/symplectic_autoencoder.jl @@ -95,23 +95,10 @@ function SymplecticAutoencoder( @assert n_encoder_blocks ≤ full_dim - reduced_dim "The number of encoder blocks is too big!" @assert n_decoder_blocks ≤ full_dim - reduced_dim "The number of decoder blocks is too big!" - if encoder_init_q && decoder_init_q - SymplecticAutoencoder{:EncoderInitQ, :DecoderInitQ, typeof(activation)}( - full_dim, reduced_dim, n_encoder_layers, n_encoder_blocks, n_decoder_layers, - n_decoder_blocks, n_decoder_output_layers, sympnet_upscale, activation) - elseif encoder_init_q && !decoder_init_q - SymplecticAutoencoder{:EncoderInitQ, :DecoderInitP, typeof(activation)}( - full_dim, reduced_dim, n_encoder_layers, n_encoder_blocks, n_decoder_layers, - n_decoder_blocks, n_decoder_output_layers, sympnet_upscale, activation) - elseif !encoder_init_q && decoder_init_q - SymplecticAutoencoder{:EncoderInitP, :DecoderInitQ, typeof(activation)}( - full_dim, reduced_dim, n_encoder_layers, n_encoder_blocks, n_decoder_layers, - n_decoder_blocks, n_decoder_output_layers, sympnet_upscale, activation) - elseif !encoder_init_q && !decoder_init_q - SymplecticAutoencoder{:EncoderInitP, :DecoderInitP, typeof(activation)}( - full_dim, reduced_dim, n_encoder_layers, n_encoder_blocks, n_decoder_layers, - n_decoder_blocks, n_decoder_output_layers, sympnet_upscale, activation) - end + SymplecticAutoencoder{encoder_init_q ? :EncoderInitQ : :EncoderInitP, + decoder_init_q ? :DecoderInitQ : :DecoderInitP, typeof(activation)}( + full_dim, reduced_dim, n_encoder_layers, n_encoder_blocks, n_decoder_layers, + n_decoder_blocks, n_decoder_output_layers, sympnet_upscale, activation) end # The intermediate steps from the full dimension to the reduced dimension, in ascending order. Only diff --git a/src/architectures/symplectic_transformer.jl b/src/architectures/symplectic_transformer.jl index 80e8b6a17..e87688c17 100644 --- a/src/architectures/symplectic_transformer.jl +++ b/src/architectures/symplectic_transformer.jl @@ -66,34 +66,18 @@ struct SymplecticTransformer{SAT, AAT, Upscaling} <: TransformerIntegrator end end -function _make_block_for_initialization(layers::Tuple, arch::SymplecticTransformer{AT}, - is_upper_criterion::Function) where {AT} - for i in 1:arch.n_sympnet - layers = (layers..., - if is_upper_criterion(i) - GradientLayerQ(arch.transformer_dim, arch.upscaling_dimension, arch.sympnet_activation) - else - GradientLayerP(arch.transformer_dim, arch.upscaling_dimension, arch.sympnet_activation) - end - ) - end - - layers -end - function create_layers_for_transformer_dimension_equal_system_dimension(arch::SymplecticTransformer) - # if `isodd(i)` and `arch.init_upper==true`, then do `GradientLayerQ`. - is_upper_criterion = arch.init_upper ? isodd : iseven + gradient_layers = _gradient_layers(arch.transformer_dim, arch.upscaling_dimension, + arch.sympnet_activation, arch.n_sympnet, arch.init_upper) layers = () for _ in 1:arch.L layers = (layers..., SymplecticAttentionQ(arch.transformer_dim; symmetric = arch.symmetric, - activation = arch.attention_activation)) - layers = _make_block_for_initialization(layers, arch, is_upper_criterion) - layers = (layers..., + activation = arch.attention_activation), + gradient_layers..., SymplecticAttentionP(arch.transformer_dim; symmetric = arch.symmetric, - activation = arch.attention_activation)) - layers = _make_block_for_initialization(layers, arch, is_upper_criterion) + activation = arch.attention_activation), + gradient_layers...) end layers diff --git a/src/architectures/sympnet.jl b/src/architectures/sympnet.jl index 9895b35a6..d7d6f8489 100644 --- a/src/architectures/sympnet.jl +++ b/src/architectures/sympnet.jl @@ -64,17 +64,10 @@ struct LASympNet{AT, InitUpperLinear, InitUpperAct} <: init_upper_act = la_init_upper_act_default) new{typeof(activation), init_upper_linear, init_upper_act}(dim, min(depth, 5), nhidden, activation) end - - function LASympNet(dl::DataLoader; depth = la_depth_default, - nhidden = la_nhidden_default, - activation = la_activation_default, - init_upper_linear = la_init_upper_linear_default, - init_upper_act = la_init_upper_act_default) - new{typeof(activation), init_upper_linear, init_upper_act}( - dl.input_dim, min(depth, 5), nhidden, activation) - end end +LASympNet(dl::DataLoader; kwargs...) = LASympNet(dl.input_dim; kwargs...) + @inline AbstractNeuralNetworks.dim(arch::SympNet) = arch.dim const g_n_layers_default = 2 @@ -103,7 +96,7 @@ struct GSympNet{AT} <: SympNet{AT} dim::Int upscaling_dimension::Int n_layers::Int - act::AT + activation::AT init_upper::Bool function GSympNet(dim; upscaling_dimension = 2 * dim, @@ -112,103 +105,41 @@ struct GSympNet{AT} <: SympNet{AT} init_upper = g_init_upper_default) new{typeof(activation)}(dim, upscaling_dimension, n_layers, activation, init_upper) end +end - function GSympNet(dl::DataLoader; upscaling_dimension = 2 * dl.input_dim, - n_layers = g_n_layers_default, - activation = g_activation_default, - init_upper = g_init_upper_default) - new{typeof(activation)}( - dl.input_dim, upscaling_dimension, n_layers, activation, init_upper) - end +GSympNet(dl::DataLoader; kwargs...) = GSympNet(dl.input_dim; kwargs...) + +# `n` gradient layers that alternate between `GradientLayerQ` and `GradientLayerP`, starting with +# `GradientLayerQ` if `init_upper` holds. +function _gradient_layers( + dim::Integer, upscaling_dimension::Integer, activation, n::Integer, + init_upper::Bool) + ntuple( + i -> isodd(i) == init_upper ? GradientLayerQ(dim, upscaling_dimension, activation) : + GradientLayerP(dim, upscaling_dimension, activation), + n) end function Chain(arch::GSympNet) - layers = () - is_upper_criterion = arch.init_upper ? isodd : iseven - for i in 1:arch.n_layers - layers = (layers..., - if is_upper_criterion(i) - GradientLayerQ(arch.dim, arch.upscaling_dimension, arch.act) - else - GradientLayerP(arch.dim, arch.upscaling_dimension, arch.act) - end - ) - end - Chain(layers...) + Chain(_gradient_layers( + arch.dim, arch.upscaling_dimension, arch.activation, arch.n_layers, arch.init_upper)...) end -function Chain(arch::LASympNet{AT, true, false}) where {AT} +# `nhidden` blocks of `depth` linear layers, a bias and two activation layers, then `depth` linear +# layers and a bias. The linear layers alternate and start with `LinearLayerQ` if +# `InitUpperLinear`; the activation layers are `ActivationLayerQ` first if `InitUpperAct`. +function Chain(arch::LASympNet{ + <:Any, InitUpperLinear, InitUpperAct}) where {InitUpperLinear, InitUpperAct} + linear_layers = ntuple( + j -> isodd(j) == InitUpperLinear ? LinearLayerQ(arch.dim) : LinearLayerP(arch.dim), + arch.depth) + q_activation = ActivationLayerQ(arch.dim, arch.activation) + p_activation = ActivationLayerP(arch.dim, arch.activation) + activation_layers = InitUpperAct ? (q_activation, p_activation) : + (p_activation, q_activation) layers = () for _ in 1:arch.nhidden - for j in 1:(arch.depth) - layers = isodd(j) ? (layers..., LinearLayerQ(arch.dim)) : - (layers..., LinearLayerP(arch.dim)) - end - layers = (layers..., BiasLayer(arch.dim)) - layers = (layers..., ActivationLayerP(arch.dim, arch.activation)) - layers = (layers..., ActivationLayerQ(arch.dim, arch.activation)) - end - for j in 1:(arch.depth) - layers = isodd(j) ? (layers..., LinearLayerQ(arch.dim)) : - (layers..., LinearLayerP(arch.dim)) - end - layers = (layers..., BiasLayer(arch.dim)) - Chain(layers...) -end - -function Chain(arch::LASympNet{AT, false, true}) where {AT} - layers = () - for i in 1:arch.nhidden - for j in 1:arch.depth - layers = isodd(j) ? (layers..., LinearLayerP(arch.dim)) : - (layers..., LinearLayerQ(arch.dim)) - end - layers = (layers..., BiasLayer(arch.dim)) - layers = (layers..., ActivationLayerQ(arch.dim, arch.activation)) - layers = (layers..., ActivationLayerP(arch.dim, arch.activation)) - end - for j in 1:(arch.depth) - layers = isodd(j) ? (layers..., LinearLayerP(arch.dim)) : - (layers..., LinearLayerQ(arch.dim)) - end - layers = (layers..., BiasLayer(arch.dim)) - Chain(layers...) -end - -function Chain(arch::LASympNet{AT, false, false}) where {AT} - layers = () - for i in 1:arch.nhidden - for j in 1:arch.depth - layers = isodd(j) ? (layers..., LinearLayerP(arch.dim)) : - (layers..., LinearLayerQ(arch.dim)) - end - layers = (layers..., BiasLayer(arch.dim)) - layers = (layers..., ActivationLayerP(arch.dim, arch.activation)) - layers = (layers..., ActivationLayerQ(arch.dim, arch.activation)) - end - for j in 1:(arch.depth) - layers = isodd(j) ? (layers..., LinearLayerP(arch.dim)) : - (layers..., LinearLayerQ(arch.dim)) - end - layers = (layers..., BiasLayer(arch.dim)) - Chain(layers...) -end - -function Chain(arch::LASympNet{AT, true, true}) where {AT} - layers = () - for i in 1:arch.nhidden - for j in 1:arch.depth - layers = isodd(j) ? (layers..., LinearLayerQ(arch.dim)) : - (layers..., LinearLayerP(arch.dim)) - end - layers = (layers..., BiasLayer(arch.dim)) - layers = (layers..., ActivationLayerQ(arch.dim, arch.activation)) - layers = (layers..., ActivationLayerP(arch.dim, arch.activation)) - end - for j in 1:(arch.depth) - layers = isodd(j) ? (layers..., LinearLayerQ(arch.dim)) : - (layers..., LinearLayerP(arch.dim)) + layers = (layers..., linear_layers..., BiasLayer(arch.dim), activation_layers...) end - layers = (layers..., BiasLayer(arch.dim)) - Chain(layers...) + Chain(layers..., linear_layers..., BiasLayer(arch.dim)) end diff --git a/src/architectures/transformer_neural_network.jl b/src/architectures/transformer_neural_network.jl index becb27f10..3fe5de48e 100644 --- a/src/architectures/transformer_neural_network.jl +++ b/src/architectures/transformer_neural_network.jl @@ -28,16 +28,17 @@ struct ClassificationTransformer{AT} <: Architecture new{AT}(transformer_dim, classification_dim, n_heads, n_layers, σ, Stiefel, add_connection) end - function ClassificationTransformer(dl::DataLoader{T, BT, CT}; - n_heads::Int = 7, - L::Int = 16, - activation::AT = softmax, - Stiefel::Bool = true, - add_connection::Bool = true) where { - T, T1, BT <: AbstractArray{T, 3}, CT <: AbstractArray{T1, 3}, AT} - new{AT}( - dl.input_dim, dl.output_dim, n_heads, L, activation, Stiefel, add_connection) - end +end + +function ClassificationTransformer(dl::DataLoader{T, BT, CT}; + n_heads::Int = 7, + L::Int = 16, + activation = softmax, + Stiefel::Bool = true, + add_connection::Bool = true) where { + T, T1, BT <: AbstractArray{T, 3}, CT <: AbstractArray{T1, 3}} + ClassificationTransformer( + dl.input_dim, dl.output_dim, n_heads, L, activation, Stiefel, add_connection) end function Chain(arch::ClassificationTransformer) diff --git a/src/layers/linear_symplectic_attention.jl b/src/layers/linear_symplectic_attention.jl index 9340c273c..d3f3e0151 100644 --- a/src/layers/linear_symplectic_attention.jl +++ b/src/layers/linear_symplectic_attention.jl @@ -72,19 +72,16 @@ parameterlength(l::LinearSymplecticAttention) = (l.seq_length + 1) * l.seq_lengt function initialparameters( rng::AbstractRNG, initializer::AbstractNeuralNetworks.Initializer, - l::LinearSymplecticAttention, backend::KernelAbstractions.Backend, T::Type) + l::LinearSymplecticAttention, backend::KernelAbstractions.Backend, ::Type{T}) where {T} S = KernelAbstractions.allocate(backend, T, parameterlength(l)) initializer(rng, S) (A = SymmetricMatrix(S, l.seq_length),) 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 + _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 + _mul(z.q, ps.A)) +function (::LinearSymplecticAttention{M, N, LayerType})( + z::NamedTuple{(:q, :p), Tuple{AT, AT}}, ps::NamedTuple) where {M, N, LayerType, AT} + _shift(Val(LayerType), z, x -> _mul(x, ps.A)) end function (d::LinearSymplecticAttention)(z::AbstractArray, ps::NamedTuple) diff --git a/src/layers/multi_head_attention.jl b/src/layers/multi_head_attention.jl index fc845fc8e..a7a492808 100644 --- a/src/layers/multi_head_attention.jl +++ b/src/layers/multi_head_attention.jl @@ -43,77 +43,50 @@ function parameterlength(d::MultiHeadAttention{M, M, true}) where {M} 3*M^2 - (3*M*(M + d.n_heads)) ÷ (2*d.n_heads) end -function initialparameters( - rng::AbstractRNG, initializer::AbstractNeuralNetworks.Initializer, - d::MultiHeadAttention{M, M, false}, - backend::KernelAbstractions.Backend, T::Type) where {M} - # number of "hidden" dimension (dimension of projection) - Dₕ = M ÷ d.n_heads - # projections for queries, keys and values. - PQ = NamedTuple() - PK = NamedTuple() - PV = NamedTuple() - - for head in 1:d.n_heads - key = Symbol("head_"*string(head)) - - PQ_weight = KernelAbstractions.allocate(backend, T, M, Dₕ) - PK_weight = KernelAbstractions.allocate(backend, T, M, Dₕ) - PV_weight = KernelAbstractions.allocate(backend, T, M, Dₕ) - initializer(rng, PQ_weight) - initializer(rng, PK_weight) - initializer(rng, PV_weight) - - PQ = merge(PQ, - NamedTuple{(key,)}((PQ_weight,)) - ) - PK = merge(PK, - NamedTuple{(key,)}((PK_weight,)) - ) - PV = merge(PV, - NamedTuple{(key,)}((PV_weight,)) - ) - end - (PQ = PQ, PK = PK, PV = PV) +# The projection of one head, `M × Dₕ`: a plain matrix, or a point of the Stiefel manifold. +function _head_weight(rng, initializer, ::Val{false}, M, Dₕ, backend, ::Type{T}) where {T} + weight = KernelAbstractions.allocate(backend, T, M, Dₕ) + initializer(rng, weight) + weight +end +function _head_weight(rng, initializer, ::Val{true}, M, Dₕ, backend, ::Type{T}) where {T} + initialparameters(rng, initializer, StiefelLayer{M, Dₕ}(), backend, T).weight end function initialparameters( rng::AbstractRNG, initializer::AbstractNeuralNetworks.Initializer, - d::MultiHeadAttention{M, M, true}, - backend::KernelAbstractions.Backend, ::Type{T}) where {M, T} - # number of "hidden" dimension (dimension of projection) + d::MultiHeadAttention{M, M, Stiefel}, + backend::KernelAbstractions.Backend, ::Type{T}) where {M, Stiefel, T} + # the dimension of the projection of one head Dₕ = M ÷ d.n_heads - # projections for queries, keys and vectors. - PQ = NamedTuple() - PK = NamedTuple() - PV = NamedTuple() - - for head in 1:d.n_heads - key = Symbol("head_"*string(head)) - - PQ = merge(PQ, - NamedTuple{(key,)}(values(initialparameters( - rng, initializer, StiefelLayer{M, Dₕ}(), backend, T))) - ) - PK = merge(PK, - NamedTuple{(key,)}(values(initialparameters( - rng, initializer, StiefelLayer{M, Dₕ}(), backend, T))) - ) - PV = merge(PV, - NamedTuple{(key,)}(values(initialparameters( - rng, initializer, StiefelLayer{M, Dₕ}(), backend, T))) - ) - end - (PQ = PQ, PK = PK, PV = PV) + # The draws go head by head, and in each head to the query, the key and the value projection. + heads = ntuple( + _ -> (PQ = _head_weight(rng, initializer, Val(Stiefel), M, Dₕ, backend, T), + PK = _head_weight(rng, initializer, Val(Stiefel), M, Dₕ, backend, T), + PV = _head_weight(rng, initializer, Val(Stiefel), M, Dₕ, backend, T)), + d.n_heads) + head_keys = ntuple(i -> Symbol("head_", i), d.n_heads) + (PQ = NamedTuple{head_keys}(map(head -> head.PQ, heads)), + PK = NamedTuple{head_keys}(map(head -> head.PK, heads)), + PV = NamedTuple{head_keys}(map(head -> head.PV, heads))) +end + +# The keys `(:head_1, …, :head_N)` of `N` heads, as a constant: the generated function builds the +# symbols once for each `N`, when it compiles, and not on each call. +@generated function _head_keys(::Val{N}) where {N} + Expr(:tuple, (QuoteNode(Symbol("head_", i)) for i in 1:N)...) end +# The forward methods find the projections of head `i` by the name `head_i`, whatever the order of +# the fields of `ps.PQ`, `ps.PK` and `ps.PV`. function compute_output_of_mha(d::MultiHeadAttention{M, M}, x::AbstractMatrix{T}, ps::NamedTuple) where { M, T} dim = size(x, 1) @assert dim == M + head_keys = _head_keys(Val(length(ps.PQ))) head_outputs = map(1:(d.n_heads)) do i - key = Symbol("head_", i) + key = head_keys[i] _mul(ps.PV[key]', x) * d.activation(_mul(ps.PQ[key]', x)' * _mul(ps.PK[key]', x) / T(sqrt(dim))) end @@ -132,8 +105,9 @@ function compute_output_of_mha( @assert dim == M # the result of a single head attention block, one per head + head_keys = _head_keys(Val(length(ps.PQ))) head_outputs = map(1:(d.n_heads)) do i - key = Symbol("head_", i) + key = head_keys[i] Q_tensor = _mul(ps.PQ[key]', x) K_tensor = _mul(ps.PK[key]', x) V_tensor = _mul(ps.PV[key]', x) diff --git a/src/layers/positional_encoding.jl b/src/layers/positional_encoding.jl index 42cc76906..802dac6d3 100644 --- a/src/layers/positional_encoding.jl +++ b/src/layers/positional_encoding.jl @@ -35,25 +35,39 @@ positional_encoding(Float32, 4, 3)[:, 1] ``` """ function positional_encoding(::Type{T}, dim::Integer, seq_length::Integer) where {T} - P = Matrix{T}(undef, dim, seq_length) - for i in 1:dim - # Row pairs `(1,2)`, `(3,4)`, … share the frequency `λ^j`, so the pair index is the row - # index shifted before the halving. `i ÷ 2` would pair `(2,3)`, `(4,5)`, … instead, and - # leave row 1 unpaired. - j = (i - 1) ÷ 2 - λʲ = T(10)^(T(-8) * T(j) / T(dim)) - for pos in 1:seq_length - θ = T(pos - 1) * λʲ - P[i, pos] = isodd(i) ? sin(θ) : cos(θ) - end - end - P + _positional_encoding!(Matrix{T}(undef, dim, seq_length)) end -# The matrix is a constant: its arguments are a type and two lengths, and none of them is a -# differentiable quantity. Without this declaration Zygote traces into the loop above and refuses -# the `setindex!`, so a network carrying a `PositionalEncoding` cannot be trained at all. +# The entry in row `i` and column `pos` of the `dim`-row encoding. +function _positional_encoding_entry(::Type{T}, dim::Integer, i::Integer, pos::Integer) where {T} + # Row pairs `(1,2)`, `(3,4)`, … share the frequency `λ^j`, so the pair index is the row index + # shifted before the halving. `i ÷ 2` would pair `(2,3)`, `(4,5)`, … instead, and leave row 1 + # unpaired. + j = (i - 1) ÷ 2 + λʲ = T(10)^(T(-8) * T(j) / T(dim)) + θ = T(pos - 1) * λʲ + isodd(i) ? sin(θ) : cos(θ) +end + +# Fill `P` with the encoding of its size, by one broadcast over its indices, so that `P` may live +# on any backend. +function _positional_encoding!(P::AbstractMatrix{T}) where {T} + dim = size(P, 1) + entry(i, pos) = _positional_encoding_entry(T, dim, i, pos) + P .= entry.(axes(P, 1), transpose(axes(P, 2))) +end + +# The encoding for the input `x` of a layer of dimension `M`, allocated on the backend of `x`. +function _positional_encoding_like(x::AbstractArray, M::Integer) + _positional_encoding!(similar(x, eltype(x), M, size(x, 2))) +end + +# The matrix is a constant: it depends on its arguments only through a type, a backend and lengths, +# and none of them is a differentiable quantity. Without these declarations Zygote traces into the +# broadcast that fills the matrix and refuses the mutation, so a network carrying a +# `PositionalEncoding` cannot be trained at all. ChainRulesCore.@non_differentiable positional_encoding(::Any, ::Any, ::Any) +ChainRulesCore.@non_differentiable _positional_encoding_like(::Any, ::Any) @doc raw""" PositionalEncoding(dim) @@ -63,16 +77,8 @@ A layer that adds the sinusoidal positional encoding of [vaswani2017attention](@ It has no parameters. The sequence length is read from the input's second axis at every call rather than fixed at construction, because a transformer here takes the sequence length from its data and not from its architecture — the same network is applied to trajectories of different lengths. The -encoding matrix is therefore built per call, which costs one `dim × seq_length` allocation, and a -second for the sum. - -!!! warning "CPU only" - [`positional_encoding`](@ref) builds a `Matrix`, so this layer adds a host array to whatever it - is given. Every other layer here allocates through the backend — `KernelAbstractions.allocate`, - or `similar(x, …)` in the forward pass — and this one does not yet. On a GPU array this is not a - slowdown but a failure: the broadcast fails to compile, because the host array cannot be read - from a kernel. A network carrying this layer is therefore a CPU network. Nothing in the test - suite would catch that, because the GPU tests in `test/metal/` do not cover this layer. +encoding matrix is therefore built per call, on the backend of the input, which costs one +`dim × seq_length` allocation, and a second for the sum. Its Jacobian is the identity, since it adds a constant, so it composes with the structure-preserving architectures without changing what they preserve: a [`LinearSymplecticTransformer`](@ref) with this @@ -117,7 +123,7 @@ parameterlength(::PositionalEncoding) = 0 # `MultiHeadAttention` rejects anyway. function (::PositionalEncoding{M, M})(x::Union{AbstractMatrix, AbstractArray{<:Any, 3}}, ::NamedTuple) where {M} - x .+ positional_encoding(eltype(x), M, size(x, 2)) + x .+ _positional_encoding_like(x, M) end # A `(q, p)` input is stacked, encoded at the full dimension `M`, and split again — which is what diff --git a/src/layers/psd_like_layer.jl b/src/layers/psd_like_layer.jl index e38de984b..456839d52 100644 --- a/src/layers/psd_like_layer.jl +++ b/src/layers/psd_like_layer.jl @@ -27,7 +27,7 @@ end function initialparameters( rng::AbstractRNG, initializer::AbstractNeuralNetworks.Initializer, - ::PSDLayer{M, N}, backend::KernelAbstractions.Backend, T::Type) where {M, N} + ::PSDLayer{M, N}, backend::KernelAbstractions.Backend, ::Type{T}) where {M, N, T} weight = N > M ? KernelAbstractions.allocate(backend, T, N ÷ 2, M ÷ 2) : KernelAbstractions.allocate(backend, T, M ÷ 2, N ÷ 2) (weight = StiefelManifold(orthonormal_columns() do diff --git a/src/layers/resnet.jl b/src/layers/resnet.jl index 50abfd600..d438f460e 100644 --- a/src/layers/resnet.jl +++ b/src/layers/resnet.jl @@ -14,41 +14,42 @@ The ResNet layer takes the following arguments: The following is a keyword argument: - `use_bias::Bool = true`: This determines whether a bias ``b`` is used. """ -struct ResNetLayer{M, N, use_bias, F1} <: AbstractExplicitLayer{M, N} +struct ResNetLayer{M, N, bias, F1} <: AbstractExplicitLayer{M, N} activation::F1 end function ResNetLayer(dim::Integer, activation = identity; use_bias::Bool = true) - ResNetLayer{dim, dim, use_bias, typeof(activation)}(activation) + ResNetLayer{dim, dim, use_bias ? :bias : :no_bias, typeof(activation)}(activation) +end + +# The bias of a layer whose bias parameter is `:bias` or `:no_bias`: a vector of length `M` filled +# by `init_bias`, or an empty `NamedTuple`. A layer merges it with its weight, which it draws first. +function _bias_parameters(::Val{:no_bias}, rng, init_bias, backend, ::Type{T}, M) where {T} + NamedTuple() +end +function _bias_parameters(::Val{:bias}, rng, init_bias, backend, ::Type{T}, M) where {T} + bias = KernelAbstractions.allocate(backend, T, M) + init_bias(rng, bias) + (bias = bias,) end function initialparameters( rng::Random.AbstractRNG, init_weight::AbstractNeuralNetworks.Initializer, - ::ResNetLayer{M, M, use_bias}, backend::KernelAbstractions.Backend, - ::Type{T}; init_bias = ZeroInitializer()) where {M, use_bias, T} - if use_bias - weight = KernelAbstractions.allocate(backend, T, M, M) - bias = KernelAbstractions.allocate(backend, T, M) - init_weight(rng, weight) - init_bias(rng, bias) - return (weight = weight, - bias = bias) - else - weight = KernelAbstractions.allocate(backend, T, M, M) - init_weight(rng, weight) - return (weight = weight,) - end -end - -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::AbstractArray, ps::NamedTuple) where {M} + ::ResNetLayer{M, M, bias}, backend::KernelAbstractions.Backend, + ::Type{T}; init_bias = ZeroInitializer()) where {M, bias, T} + weight = KernelAbstractions.allocate(backend, T, M, M) + init_weight(rng, weight) + merge((weight = weight,), _bias_parameters(Val(bias), rng, init_bias, backend, T, M)) +end + +parameterlength(::ResNetLayer{M, M, :bias}) where {M} = M * (M + 1) +parameterlength(::ResNetLayer{M, M, :no_bias}) where {M} = M * M + +function (d::ResNetLayer{M, M, :bias})(x::AbstractArray, ps::NamedTuple) where {M} x + d.activation.(_mul(ps.weight, x) .+ ps.bias) end -function (d::ResNetLayer{M, M, false})(x::AbstractArray, ps::NamedTuple) where {M} +function (d::ResNetLayer{M, M, :no_bias})(x::AbstractArray, ps::NamedTuple) where {M} x + d.activation.(_mul(ps.weight, x)) end diff --git a/src/layers/symplectic_attention.jl b/src/layers/symplectic_attention.jl index a84502388..c94125c4e 100644 --- a/src/layers/symplectic_attention.jl +++ b/src/layers/symplectic_attention.jl @@ -74,16 +74,14 @@ const SymplecticAttentionP{M, N, Symmetric, AT} = SymplecticAttention{ function SymplecticAttentionQ(M::Integer; symmetric::Bool = sa_symmetric_default, activation::AbstractSoftmax = sa_activation_default) @assert iseven(M) "Dimension must be even!" - AT = typeof(activation) - !symmetric ? SymplecticAttention{M, M, :Q, :arbitrary, AT}(activation) : - SymplecticAttention{M, M, :Q, :symmetric, AT}(activation) + SymplecticAttention{M, M, :Q, symmetric ? :symmetric : :arbitrary, typeof(activation)}( + activation) end function SymplecticAttentionP(M::Integer; symmetric::Bool = sa_symmetric_default, activation::AbstractSoftmax = sa_activation_default) @assert iseven(M) "Dimension must be even!" - AT = typeof(activation) - !symmetric ? SymplecticAttention{M, M, :P, :arbitrary, AT}(activation) : - SymplecticAttention{M, M, :P, :symmetric, AT}(activation) + SymplecticAttention{M, M, :P, symmetric ? :symmetric : :arbitrary, typeof(activation)}( + activation) end function parameterlength(::SymplecticAttention{ @@ -102,7 +100,7 @@ function initialparameters(rng::AbstractRNG, initializer::AbstractNeuralNetworks.Initializer, l::SymplecticAttention{M, M, LayerType, :symmetric}, backend::KernelAbstractions.Backend, - T::Type) where {M, LayerType} + ::Type{T}) where {M, LayerType, T} S = KernelAbstractions.allocate(backend, T, parameterlength(l)) initializer(rng, S) (A = SymmetricMatrix(S, M ÷ 2),) @@ -112,20 +110,20 @@ function initialparameters(rng::AbstractRNG, initializer::AbstractNeuralNetworks.Initializer, ::SymplecticAttention{M, M, LayerType, :arbitrary}, backend::KernelAbstractions.Backend, - T::Type) where {M, LayerType} + ::Type{T}) where {M, LayerType, T} A = KernelAbstractions.allocate(backend, T, M ÷ 2, M ÷ 2) initializer(rng, A) (A = A,) end -function (d::SymplecticAttentionQ{M, M, :arbitrary})( - z::NamedTuple{(:q, :p), Tuple{AT, AT}}, ps::NamedTuple) where {AT, M} - PAP = _mul(_mul(_custom_transpose(z.p), ps.A), z.p) - σPAP = d.activation(PAP) - ( - q = z.q + (_mul(_mul(ps.A, z.p), _custom_transpose(σPAP)) + - _mul(_mul(ps.A', z.p), σPAP)), - p = z.p) +function (d::SymplecticAttention{M, M, LayerType, :arbitrary})( + z::NamedTuple{(:q, :p), Tuple{AT, AT}}, ps::NamedTuple) where {M, LayerType, AT} + _shift(Val(LayerType), z, + function (x) + XAX = _mul(_mul(_custom_transpose(x), ps.A), x) + σXAX = d.activation(XAX) + _mul(_mul(ps.A, x), _custom_transpose(σXAX)) + _mul(_mul(ps.A', x), σXAX) + end) end # `A = ps.A` once, and then `A` twice: the local binding is what keeps the `SymmetricMatrix` @@ -133,30 +131,11 @@ end # to a plain `Matrix` -- `test/parameters/symplectic_attention_network_parameters_gradient.jl` # asserts both halves of that, and the last testset there covers this method. The binding costs no # allocation; it is the number of accesses that decides the structure. -function (d::SymplecticAttentionQ{M, M, :symmetric})( - z::NamedTuple{(:q, :p), Tuple{AT, AT}}, ps::NamedTuple) where {AT, M} +function (d::SymplecticAttention{M, M, LayerType, :symmetric})( + z::NamedTuple{(:q, :p), Tuple{AT, AT}}, ps::NamedTuple) where {M, LayerType, AT} A = ps.A - PAP = _mul(_mul(_custom_transpose(z.p), A), z.p) - σPAP = d.activation(PAP) - (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 = _mul(_mul(_custom_transpose(z.q), ps.A), z.q) - σQAQ = d.activation(QAQ) - (q = z.q, - 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 = _mul(_mul(_custom_transpose(z.q), A), z.q) - σQAQ = d.activation(QAQ) - (q = z.q, p = z.p + _mul(_mul(A, z.q), 2 * σQAQ)) + _shift(Val(LayerType), z, + x -> _mul(_mul(A, x), 2 * d.activation(_mul(_mul(_custom_transpose(x), A), x)))) end function (d::SymplecticAttention)(z::AbstractArray, ps::NamedTuple) diff --git a/src/layers/sympnets.jl b/src/layers/sympnets.jl index a95ff29c1..cc469f0b8 100644 --- a/src/layers/sympnets.jl +++ b/src/layers/sympnets.jl @@ -15,7 +15,7 @@ abstract type SympNetLayer{M, N} <: AbstractExplicitLayer{M, N} end See the docstrings for the constructors [`GradientLayerQ`](@ref) and [`GradientLayerP`](@ref). """ struct GradientLayer{M, N, TA, C} <: SympNetLayer{M, N} - second_dim::Integer + second_dim::Int activation::TA end @@ -83,7 +83,7 @@ This is equivalent to a left multiplication by the matrix: ``` where ``A`` is a [`SymmetricMatrix`](@extref GeometricOptimizers GeometricOptimizers.SymmetricMatrix). """ -const LinearLayerQ{M, N, TA} = LinearLayer{M, N, :Q} +const LinearLayerQ{M, N} = LinearLayer{M, N, :Q} @doc raw""" LinearLayerP(n) @@ -99,7 +99,7 @@ A & \mathbb{I} ``` where ``A`` is a [`SymmetricMatrix`](@extref GeometricOptimizers GeometricOptimizers.SymmetricMatrix). """ -const LinearLayerP{M, N, TA} = LinearLayer{M, N, :P} +const LinearLayerP{M, N} = LinearLayer{M, N, :P} @doc raw""" ActivationLayer <: SympNetLayer @@ -150,28 +150,30 @@ This can be recovered from [`GradientLayerP`](@ref) by setting ``M`` equal to `n """ const ActivationLayerP{M, N, TA} = ActivationLayer{M, N, TA, :P} +# Each constructor converts the dimension to an `Int`, so that an `Int32` dimension names the same +# layer as an `Int` one. function GradientLayerQ(M, upscaling_dimension, activation) - GradientLayer{M, M, typeof(activation), :Q}(upscaling_dimension, activation) + GradientLayer{Int(M), Int(M), typeof(activation), :Q}(upscaling_dimension, activation) end function GradientLayerP(M, upscaling_dimension, activation) - GradientLayer{M, M, typeof(activation), :P}(upscaling_dimension, activation) + GradientLayer{Int(M), Int(M), typeof(activation), :P}(upscaling_dimension, activation) end function ActivationLayerQ(M, activation) - ActivationLayerQ{M, M, typeof(activation)}(activation) + ActivationLayerQ{Int(M), Int(M), typeof(activation)}(activation) end function ActivationLayerP(M, activation) - ActivationLayerP{M, M, typeof(activation)}(activation) + ActivationLayerP{Int(M), Int(M), typeof(activation)}(activation) end function LinearLayerQ(M) - LinearLayerQ{M, M}() + LinearLayerQ{Int(M), Int(M)}() end function LinearLayerP(M) - LinearLayerP{M, M}() + LinearLayerP{Int(M), Int(M)}() end function initialparameters( @@ -215,39 +217,25 @@ function parameterlength(::LinearLayer{M, M}) where {M} (M ÷ 2) * (M ÷ 2 + 1) ÷ 2 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 + 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 + 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 + _mul(ps.weight', - ps.scale .* d.activation.(_mul(ps.weight, x.p) .+ ps.bias)), - p = x.p) -end +# The update of a layer that changes one half of `z`: a `:Q` layer adds `f(z.p)` to `z.q`, a `:P` +# layer adds `f(z.q)` to `z.p`. +_shift(::Val{:Q}, z::NamedTuple, f) = (q = z.q + f(z.p), p = z.p) +_shift(::Val{:P}, z::NamedTuple, f) = (q = z.q, p = z.p + f(z.q)) -@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 + _mul(ps.weight', - ps.scale .* d.activation.(_mul(ps.weight, x.q) .+ ps.bias))) +@inline function (d::ActivationLayer{M, M, TA, C})(x::NamedTuple, ps) where {M, TA, C} + size(x.q, 1) == size(x.p, 1) == M ÷ 2 || error("Dimension mismatch.") + _shift(Val(C), x, y -> ps.scale .* d.activation.(y)) end -@inline function (d::LinearLayerQ{M, M})(x::NamedTuple, ps) where {M} - size(x.q, 1) == M÷2 || error("Dimension mismatch.") - (q = x.q + _mul(ps.weight, x.p), p = x.p) +@inline function (d::GradientLayer{M, M, TA, C})(x::NamedTuple, ps) where {M, TA, C} + size(x.q, 1) == size(x.p, 1) == M ÷ 2 || error("Dimension mismatch.") + _shift(Val(C), x, + y -> _mul(ps.weight', ps.scale .* d.activation.(_mul(ps.weight, y) .+ ps.bias))) 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 + _mul(ps.weight, x.q)) +@inline function (d::LinearLayer{M, M, C})(x::NamedTuple, ps) where {M, C} + size(x.q, 1) == size(x.p, 1) == M ÷ 2 || error("Dimension mismatch.") + _shift(Val(C), x, y -> _mul(ps.weight, y)) end # For the wrappers whose input to a SympNet layer is an `AbstractArray` rather than the `NamedTuple` diff --git a/src/layers/transformer.jl b/src/layers/transformer.jl index 0a6be5edf..4e52354c8 100644 --- a/src/layers/transformer.jl +++ b/src/layers/transformer.jl @@ -24,8 +24,7 @@ Make an instance of the Transformer with `n_heads` for dimension `dim` and `L` b raw""": if a [`PositionalEncoding`](@ref) should be put in front of the chain. Attention is permutation-equivariant on its own, so without this the network cannot tell one ordering of a sequence from another. It is off by default because the transformers here are usually applied - to phase-space trajectories, where the ordering is carried by the data. Setting it to `true` - makes the network CPU-only, for the reason given in [`PositionalEncoding`](@ref). + to phase-space trajectories, where the ordering is carried by the data. """ function Transformer(dim::Integer, n_heads::Integer, L::Integer; activation = t_activation_default, diff --git a/src/layers/volume_preserving_attention.jl b/src/layers/volume_preserving_attention.jl index 5d299591c..638049558 100644 --- a/src/layers/volume_preserving_attention.jl +++ b/src/layers/volume_preserving_attention.jl @@ -25,11 +25,7 @@ device, and in the element type of the data. """ struct VolumePreservingAttention{M, N, ScalarProductType, SL} <: AbstractExplicitLayer{M, N} function VolumePreservingAttention(dim::Int, seq_length::Int = 0; skew_sym::Bool = false) - if skew_sym - return new{dim, dim, :skew_sym, seq_length}() - else - return new{dim, dim, :arbitrary, seq_length}() - end + new{dim, dim, skew_sym ? :skew_sym : :arbitrary, seq_length}() end end @@ -49,7 +45,7 @@ end function initialparameters( rng::AbstractRNG, initializer!::AbstractNeuralNetworks.Initializer, d::VolumePreservingAttention{M, M, :skew_sym}, - backend::KernelAbstractions.Backend, T::Type) where {M} + backend::KernelAbstractions.Backend, ::Type{T}) where {M, T} V = KernelAbstractions.allocate(backend, T, parameterlength(d)) initializer!(rng, V) (A = SkewSymMatrix(V, M),) @@ -58,7 +54,7 @@ end function initialparameters( rng::AbstractRNG, initializer!::AbstractNeuralNetworks.Initializer, ::VolumePreservingAttention{M, M, :arbitrary}, - backend::KernelAbstractions.Backend, T::Type) where {M} + backend::KernelAbstractions.Backend, ::Type{T}) where {M, T} A = KernelAbstractions.allocate(backend, T, M, M) initializer!(rng, A) (A = A,) diff --git a/src/layers/volume_preserving_feedforward.jl b/src/layers/volume_preserving_feedforward.jl index 355f02c06..5cabd2b78 100644 --- a/src/layers/volume_preserving_feedforward.jl +++ b/src/layers/volume_preserving_feedforward.jl @@ -1,7 +1,7 @@ @doc raw""" VolumePreservingFeedForwardLayer <: AbstractExplicitLayer -Super-type of [`VolumePreservingLowerLayer`](@ref) and [`VolumePreservingUpperLayer`](@ref). The layers do the following: +The layer type of which [`VolumePreservingLowerLayer`](@ref) and [`VolumePreservingUpperLayer`](@ref) are the two triangles. The layers do the following: ```math x \mapsto \begin{cases} \sigma(Lx + b) & \text{where $L$ is }\mathtt{StrictlyLowerTriangular}, \\ \sigma(Ux + b) & \text{where $U$ is }\mathtt{StrictlyUpperTriangular}. \end{cases} @@ -9,7 +9,10 @@ x \mapsto \begin{cases} \sigma(Lx + b) & \text{where $L$ is }\mathtt{StrictlyLow The functor can be applied to a vector, a matrix or a tensor. The special matrices are implemented as [`StrictlyLowerTriangular`](@extref GeometricOptimizers GeometricOptimizers.StrictlyLowerTriangular) and [`StrictlyUpperTriangular`](@extref GeometricOptimizers GeometricOptimizers.StrictlyUpperTriangular). """ -abstract type VolumePreservingFeedForwardLayer{M, N, bias} <: AbstractExplicitLayer{M, N} end +struct VolumePreservingFeedForwardLayer{M, N, bias, AT, Triangle} <: + AbstractExplicitLayer{M, N} + activation::AT +end @doc raw""" VolumePreservingLowerLayer(dim) @@ -53,19 +56,15 @@ l(x, ps) The constructor can be called with the optional arguments: - `activation=tanh`: the activation function. -- `use_bias::Bool=true` (keyword argument): specifies whether a bias should be used. +- `use_bias::Bool=true` (keyword argument): specifies whether a bias should be used. """ -struct VolumePreservingLowerLayer{M, N, bias, AT} <: - VolumePreservingFeedForwardLayer{M, N, bias} - activation::AT +const VolumePreservingLowerLayer{M, N, bias, AT} = VolumePreservingFeedForwardLayer{ + M, N, bias, AT, :lower} - function VolumePreservingLowerLayer(sys_dim::Int, activation = tanh; use_bias::Bool = true) - if use_bias - return new{sys_dim, sys_dim, :bias, typeof(activation)}(activation) - else - return new{sys_dim, sys_dim, :no_bias, typeof(activation)}(activation) - end - end +function VolumePreservingLowerLayer(sys_dim::Int, activation = tanh; use_bias::Bool = true) + VolumePreservingLowerLayer{ + sys_dim, sys_dim, use_bias ? :bias : :no_bias, typeof(activation)}( + activation) end @doc raw""" @@ -110,21 +109,20 @@ l(x, ps) The constructor can be called with the optional arguments: - `activation=tanh`: the activation function. -- `use_bias::Bool=true` (keyword argument): specifies whether a bias should be used. +- `use_bias::Bool=true` (keyword argument): specifies whether a bias should be used. """ -struct VolumePreservingUpperLayer{M, N, bias, AT} <: - VolumePreservingFeedForwardLayer{M, N, bias} - activation::AT +const VolumePreservingUpperLayer{M, N, bias, AT} = VolumePreservingFeedForwardLayer{ + M, N, bias, AT, :upper} - function VolumePreservingUpperLayer(sys_dim::Int, activation = tanh; use_bias::Bool = true) - if use_bias - return new{sys_dim, sys_dim, :bias, typeof(activation)}(activation) - else - return new{sys_dim, sys_dim, :no_bias, typeof(activation)}(activation) - end - end +function VolumePreservingUpperLayer(sys_dim::Int, activation = tanh; use_bias::Bool = true) + VolumePreservingUpperLayer{ + sys_dim, sys_dim, use_bias ? :bias : :no_bias, typeof(activation)}( + activation) end +_strictly_triangular(::Val{:lower}) = StrictlyLowerTriangular +_strictly_triangular(::Val{:upper}) = StrictlyUpperTriangular + function parameterlength(::VolumePreservingFeedForwardLayer{M, M, :no_bias}) where {M} M * (M - 1) ÷ 2 end @@ -133,47 +131,13 @@ function parameterlength(::VolumePreservingFeedForwardLayer{M, M, :bias}) where end function initialparameters( - rng::AbstractRNG, init_weight!::AbstractNeuralNetworks.Initializer, - d::VolumePreservingLowerLayer{M, M, :bias}, backend::Backend, - ::Type{T}; init_bias! = ZeroInitializer()) where {M, T} - S = KernelAbstractions.allocate(backend, T, parameterlength(d) - M) - b = KernelAbstractions.allocate(backend, T, M) - init_weight!(rng, S) - init_bias!(rng, b) - - (weight = StrictlyLowerTriangular(S, M), bias = b) -end - -function initialparameters( - rng::AbstractRNG, init_weight!::AbstractNeuralNetworks.Initializer, - d::VolumePreservingUpperLayer{M, M, :bias}, backend::Backend, - ::Type{T}; init_bias! = ZeroInitializer()) where {M, T} - S = KernelAbstractions.allocate(backend, T, parameterlength(d) - M) - b = KernelAbstractions.allocate(backend, T, M) - init_weight!(rng, S) - init_bias!(rng, b) - - (weight = StrictlyUpperTriangular(S, M), bias = b) -end - -function initialparameters( - rng::AbstractRNG, init_weight!::AbstractNeuralNetworks.Initializer, - d::VolumePreservingLowerLayer{M, M, :no_bias}, backend::Backend, - ::Type{T}; init_bias! = ZeroInitializer()) where {M, T} - S = KernelAbstractions.allocate(backend, T, parameterlength(d)) - init_weight!(rng, S) - - (weight = StrictlyLowerTriangular(S, M),) -end - -function initialparameters( - rng::AbstractRNG, init_weight!::AbstractNeuralNetworks.Initializer, - d::VolumePreservingUpperLayer{M, M, :no_bias}, backend::Backend, - ::Type{T}; init_bias! = ZeroInitializer()) where {M, T} - S = KernelAbstractions.allocate(backend, T, parameterlength(d)) - init_weight!(rng, S) - - (weight = StrictlyUpperTriangular(S, M),) + rng::AbstractRNG, init_weight::AbstractNeuralNetworks.Initializer, + ::VolumePreservingFeedForwardLayer{M, M, bias, AT, Triangle}, backend::Backend, + ::Type{T}; init_bias = ZeroInitializer()) where {M, bias, AT, Triangle, T} + S = KernelAbstractions.allocate(backend, T, M * (M - 1) ÷ 2) + init_weight(rng, S) + weight = _strictly_triangular(Val(Triangle))(S, M) + merge((weight = weight,), _bias_parameters(Val(bias), rng, init_bias, backend, T, M)) end function (d::VolumePreservingFeedForwardLayer{M, M, :bias})(x::AbstractArray, ps) where {M} diff --git a/test/architectures/architecture_pins.jl b/test/architectures/architecture_pins.jl new file mode 100644 index 000000000..d644f920e --- /dev/null +++ b/test/architectures/architecture_pins.jl @@ -0,0 +1,86 @@ +# The forward pass of each architecture whose chain is built from shared code, and the outputs of +# the encoders and decoders, against the outputs that commit 4c4a8320 records for the same seed and +# input; and the layer types of `LASympNet` for its four `(init_upper_linear, init_upper_act)` +# pairs, and the encoder and decoder types of the three autoencoder families. + +using GeometricMachineLearning +using GeometricMachineLearning: AutoEncoder, NonLinearSymplecticEncoder, + NonLinearSymplecticDecoder, + UnknownSymplecticEncoder, UnknownSymplecticDecoder, + UnknownEncoder, UnknownDecoder +using Test + +include("../helpers/eltypes.jl") +include("../helpers/pins.jl") +include("../helpers/pinned_outputs.jl") + +@testset "$name ($T)" for (i, ( + name, model, dim, cols)) in enumerate(architecture_pin_cases()), + T in REAL_ELTYPES + + nn = pinned_network(model, T, architecture_seed(i)) + y = nn(pin_input(T, dim, cols)) + @test eltype(y) === T + @test pin_ulps(vec(y), ARCHITECTURE_PINS[(name, T)]) ≤ PIN_ULPS +end + +@testset "encoder and decoder of $name ($T)" for (i, ( + name, model, full, reduced)) in enumerate(autoencoder_pin_cases()), + T in REAL_ELTYPES + + nn = pinned_network(model, T, autoencoder_seed(i)) + ye = encoder(nn)(pin_input(T, full, 3)) + yd = decoder(nn)(pin_input(T, reduced, 3)) + @test eltype(ye) === T + @test eltype(yd) === T + @test pin_ulps(vec(ye), AUTOENCODER_PINS[(name, :encoder, T)]) ≤ PIN_ULPS + @test pin_ulps(vec(yd), AUTOENCODER_PINS[(name, :decoder, T)]) ≤ PIN_ULPS +end + +@testset "LASympNet layer types, init_upper_linear = $linear, init_upper_act = $act" for linear in (true, + false), + act in (true, false) + + arch = LASympNet( + 4; depth = 3, nhidden = 2, init_upper_linear = linear, init_upper_act = act) + @test [type_string(typeof(l)) for l in Chain(arch).layers] == + LASYMPNET_LAYER_TYPES[(linear, act)] +end + +# An autoencoder that defines its layers and nothing else, so that `encoder` and `decoder` reach the +# methods for a plain `AutoEncoder`. +struct PlainAE <: AutoEncoder + full_dim::Int + reduced_dim::Int + n_encoder_blocks::Int + n_decoder_blocks::Int +end + +function GeometricMachineLearning.encoder_layers_from_iteration(arch::PlainAE, ::AbstractVector{<:Integer}) + ( + Dense(arch.full_dim, arch.reduced_dim, tanh),) +end +function GeometricMachineLearning.decoder_layers_from_iteration(arch::PlainAE, ::AbstractVector{<:Integer}) + ( + Dense(arch.reduced_dim, arch.full_dim, tanh),) +end + +@testset "encoder and decoder types of each autoencoder family ($T)" for T in REAL_ELTYPES + sae = NeuralNetwork(SymplecticAutoencoder(8, 4), CPU(), T) + @test encoder(sae).architecture isa NonLinearSymplecticEncoder + @test decoder(sae).architecture isa NonLinearSymplecticDecoder + + psd = NeuralNetwork(PSDArch(8, 4), CPU(), T) + @test encoder(psd).architecture isa UnknownSymplecticEncoder + @test decoder(psd).architecture isa UnknownSymplecticDecoder + + plain = NeuralNetwork(PlainAE(6, 2, 2, 2), CPU(), T) + @test encoder(plain).architecture isa UnknownEncoder + @test decoder(plain).architecture isa UnknownDecoder + @test eltype(decoder(plain)(encoder(plain)(pin_input(T, 6, 3)))) === T + + # The encoder and the decoder each keep the block count of their own half. + uneven = NeuralNetwork(PlainAE(6, 2, 2, 3), CPU(), T) + @test encoder(uneven).architecture.n_encoder_blocks == 2 + @test decoder(uneven).architecture.n_decoder_blocks == 3 +end diff --git a/test/architectures/architecture_structure.jl b/test/architectures/architecture_structure.jl new file mode 100644 index 000000000..58b7618af --- /dev/null +++ b/test/architectures/architecture_structure.jl @@ -0,0 +1,96 @@ +# The architectures keep their keywords: a non-default activation reaches every layer that takes +# it, the activation field has one name across the SympNets, the `DataLoader` constructors of +# `LASympNet`, `GSympNet` and `ClassificationTransformer` give the architecture of the dimensions +# they read, and the chain builders that shared a body are one method. + +using GeometricMachineLearning +using GeometricMachineLearning: GradientLayer +using AbstractNeuralNetworks: Dense +using Test +import Random + +const GML = GeometricMachineLearning + +Random.seed!(1234) + +include("../helpers/eltypes.jl") + +# The methods of `f` whose definitions are in the file `file` under `src/`. +function methods_in(f, file::AbstractString) + filter( + m -> endswith(String(m.file), joinpath("src", file)), collect(methods(f))) +end + +@testset "LinearSymplecticTransformer keeps a non-default activation" begin + arch = LinearSymplecticTransformer(4, 3; activation = sin) + @test arch.activation === sin + gradient_layers = filter(l -> l isa GradientLayer, collect(Chain(arch).layers)) + @test length(gradient_layers) == 4 + @test all(l -> l.activation === sin, gradient_layers) +end + +@testset "StandardTransformerIntegrator stores and uses upscaling_activation" begin + @test hasfield(StandardTransformerIntegrator, :upscaling_activation) + arch = StandardTransformerIntegrator(4; transformer_dim = 6, n_heads = 2, + upscaling_activation = tanh) + @test arch.upscaling_activation === tanh + @test Chain(arch).layers[1].σ === Dense(4, 6, tanh).σ + + default = StandardTransformerIntegrator(4; transformer_dim = 6, n_heads = 2) + @test Chain(default).layers[1].σ === Dense(4, 6, identity).σ +end + +@testset "the SympNets name their activation field activation" begin + @test GSympNet(4; activation = sin).activation === sin + @test LASympNet(4; activation = sin).activation === sin + @test !hasfield(GSympNet, :act) +end + +@testset "a DataLoader gives the SympNet of its input dimension ($T)" for T in REAL_ELTYPES + dl = DataLoader(rand(T, 4, 10); suppress_info = true) + + la = LASympNet(dl; depth = 7, nhidden = 2, activation = sin, init_upper_linear = false, + init_upper_act = false) + @test la === LASympNet(4; depth = 7, nhidden = 2, activation = sin, + init_upper_linear = false, init_upper_act = false) + @test la.depth == 5 + + g = GSympNet(dl; n_layers = 3, activation = sin, init_upper = false) + @test g === GSympNet(4; n_layers = 3, activation = sin, init_upper = false) + @test g.upscaling_dimension == 8 +end + +@testset "a DataLoader gives the ClassificationTransformer of its dimensions ($T)" for T in REAL_ELTYPES + dl = DataLoader(rand(T, 4, 3, 5), rand(T, 2, 1, 5); suppress_info = true) + @test ClassificationTransformer(dl) === + ClassificationTransformer(4, 2, 7, 16, softmax, true, true) + @test ClassificationTransformer( + dl; n_heads = 2, L = 3, activation = tanh, Stiefel = false, add_connection = false) === + ClassificationTransformer(4, 2, 2, 3, tanh, false, false) + # The two flags differ, so a constructor that swaps them gives another architecture. + @test ClassificationTransformer(dl; Stiefel = true, add_connection = false) === + ClassificationTransformer(4, 2, 7, 16, softmax, true, false) + @test ClassificationTransformer(dl; Stiefel = false, add_connection = true) === + ClassificationTransformer(4, 2, 7, 16, softmax, false, true) +end + +@testset "one chain builder where the bodies were the same" begin + # One `Chain` method for `LASympNet`, beside the one for `GSympNet`. + @test length(methods_in(Chain, joinpath("architectures", "sympnet.jl"))) == 2 + # One builder of the alternating gradient layers for GSympNet, LST and ST. + @test !isdefined(GML, :_make_block_for_initialization) + @test length(methods(GML._gradient_layers)) == 1 + # One `encoder` and one `decoder` for every autoencoder that is not a `SymplecticAutoencoder`, + # beside the methods for an arbitrary network. + @test length(methods_in(encoder, joinpath("architectures", "autoencoder.jl"))) == 2 + @test length(methods_in(decoder, joinpath("architectures", "autoencoder.jl"))) == 2 +end + +@testset "_gradient_layers alternates from the side that init_upper names" begin + layers = GML._gradient_layers(4, 8, sin, 3, true) + @test layers == + (GradientLayerQ(4, 8, sin), GradientLayerP(4, 8, sin), GradientLayerQ(4, 8, sin)) + layers = GML._gradient_layers(4, 8, sin, 2, false) + @test layers == (GradientLayerP(4, 8, sin), GradientLayerQ(4, 8, sin)) + @test GML._gradient_layers(4, 8, sin, 0, true) == () +end diff --git a/test/architectures/forward_inference.jl b/test/architectures/forward_inference.jl new file mode 100644 index 000000000..4b94736a4 --- /dev/null +++ b/test/architectures/forward_inference.jl @@ -0,0 +1,41 @@ +# The forward pass of every architecture that the package exports infers its result type, and the +# result has the element type of the network. `GeneralizedHamiltonianArchitecture` has no chain: +# its constructor raises "GHNN still has to be implemented!". + +using GeometricMachineLearning +using Test +import Random + +include("../helpers/eltypes.jl") + +input(::Type{T}, dim, cols) where {T} = T[sin(i + 2j) for i in 1:dim, j in 1:cols] + +# `(name, architecture, input dimension, number of columns)`. +const ARCHITECTURES = ( + ("ResNet", ResNet(4, 2, tanh), 4, 3), + ("StandardTransformerIntegrator", StandardTransformerIntegrator(4; n_heads = 2), 4, 3), + # The default is one head per dimension, so this is a 32-head attention. + ("StandardTransformerIntegrator, default heads", + StandardTransformerIntegrator(32), 32, 3), + ("VolumePreservingTransformer", VolumePreservingTransformer(4, 3), 4, 3), + ("LagrangianNeuralNetwork", LagrangianNeuralNetwork(4), 4, 3), + ("LASympNet", LASympNet(4), 4, 3), + ("GSympNet", GSympNet(4), 4, 3), + ("ClassificationTransformer", + ClassificationTransformer(4, 2, 2, 1, softmax, false, true), 4, 3), + ("VolumePreservingFeedForward", VolumePreservingFeedForward(4), 4, 3), + ("SymplecticAutoencoder", SymplecticAutoencoder(8, 4), 8, 3), + ("PSDArch", PSDArch(8, 4), 8, 3), + ("StandardHamiltonianArchitecture", StandardHamiltonianArchitecture(4), 4, 3), + ("LinearSymplecticTransformer", LinearSymplecticTransformer(4, 3), 4, 3), + ("SymplecticTransformer", SymplecticTransformer(4), 4, 3)) + +@testset "the forward pass of $name infers ($T)" for (name, arch, dim, cols) in ARCHITECTURES, + T in REAL_ELTYPES + + Random.seed!(0) + nn = NeuralNetwork(arch, CPU(), T) + x = input(T, dim, cols) + y = @inferred nn(x) + @test eltype(y) === T +end diff --git a/test/helpers/pinned_outputs.jl b/test/helpers/pinned_outputs.jl new file mode 100644 index 000000000..5e059f540 --- /dev/null +++ b/test/helpers/pinned_outputs.jl @@ -0,0 +1,397 @@ +# Recorded from commit 4c4a8320 with the cases and seeds of `test/helpers/pins.jl`. + +const LAYER_PINS = Dict{Tuple{String, DataType}, Vector}( + ("GradientLayerQ", Float32) => + Float32[-0.51713926, -1.3761483, -0.9589243, -0.2794155, -0.542231, 1.1872401, + 0.6569866, 0.98935825, 0.8197795, 0.06232357, 0.4121185, -0.5440211], + ("GradientLayerQ", Float64) => + [-0.4390427803799615, -0.5799544940267445, -0.9589242746631385, + -0.27941549819892586, -0.4509930458617585, -0.883775922123758, + 0.6569865987187891, 0.9893582466233818, 0.9348843997469185, + 1.4570426417905926, 0.4121184852417566, -0.5440211108893698], + ("GradientLayerP", Float32) => + Float32[0.14112, -0.7568025, -1.0310265, -0.2679918, -0.9589243, -0.2794155, + 0.25323027, 0.8627788, 0.6569866, 0.98935825, 0.8033464, -0.43903914], + ("GradientLayerP", Float64) => + [0.1411200080598672, -0.7568024953079282, -0.30131958613169485, + -0.2731333042958512, -0.9589242746631385, -0.27941549819892586, + 0.27300215215963636, 1.710629715243495, 0.6569865987187891, + 0.9893582466233818, 0.1893898608589251, -0.9804828858697413], + ("ActivationLayerQ", Float32) => + Float32[0.2733904, -0.7225378, -0.9589243, -0.2794155, -1.0614183, -0.3746608, + 0.6569866, 0.98935825, 0.5875843, 1.0517607, 0.4121185, -0.5440211], + ("ActivationLayerQ", Float64) => + [0.23469308350274765, -0.7118650428897826, -0.9589242746631385, + -0.27941549819892586, -1.0314323505993654, -0.4043278487539732, + 0.6569865987187891, 0.9893582466233818, 0.6078888215493407, + 1.0711977407352156, 0.4121184852417566, -0.5440211108893698], + ("ActivationLayerP", Float32) => + Float32[0.14112, -0.7568025, -0.8743746, -1.0628524, -0.9589243, -0.2794155, + 0.20839894, 0.65552926, 0.6569866, 0.98935825, 0.7597209, 0.3839211], + ("ActivationLayerP", Float64) => + [0.1411200080598672, -0.7568024953079282, -0.7870965415833097, + 0.16967129124922464, -0.9589242746631385, -0.27941549819892586, + -0.25466424201071225, 1.1807178455557388, 0.6569865987187891, + 0.9893582466233818, 1.1185401570109004, -1.0759421704432448], + ("LinearLayerQ", Float32) => + Float32[1.0452667, -0.033778965, -0.9589243, -0.2794155, -2.5039887, 0.33707988, + 0.6569866, 0.98935825, 1.0387872, -0.24677044, 0.4121185, -0.5440211], + ("LinearLayerQ", Float64) => + [1.5076103585430225, 0.4922354670721971, -0.9589242746631385, + -0.27941549819892586, -2.6208254476009674, -2.2115720272073274, + 0.6569865987187891, 0.9893582466233818, 0.6736860797797595, + 1.348441938764763, 0.4121184852417566, -0.5440211108893698], + ("LinearLayerP", Float32) => + Float32[0.14112, -0.7568025, -0.12699932, 0.08234322, -0.9589243, -0.2794155, + -0.3271507, 1.9935205, 0.6569866, 0.98935825, 0.39928478, -1.7415378], + ("LinearLayerP", Float64) => + [0.1411200080598672, -0.7568024953079282, -0.7196985954738497, + -1.087516847128719, -0.9589242746631385, -0.27941549819892586, + 1.6159089487905813, 1.1729284654363297, 0.6569865987187891, + 0.9893582466233818, -0.6252121989009878, 0.11129590635387365], + ("SymplecticAttentionQ symmetric", Float32) => + Float32[-0.20285256, -1.0754645, -0.9589243, -0.2794155, -0.024502158, 0.28820854, + 0.6569866, 0.98935825, 0.7200119, 0.9644051, 0.4121185, -0.5440211], + ("SymplecticAttentionQ symmetric", Float64) => + [-0.7131136091901249, -1.2335324324898373, -0.9589242746631385, + -0.27941549819892586, -0.16590794312803236, -0.032830836804668306, + 0.6569865987187891, 0.9893582466233818, 1.0044352819476057, + 0.949139160672247, 0.4121184852417566, -0.5440211108893698], + ("SymplecticAttentionP symmetric", Float32) => + Float32[0.14112, -0.7568025, -0.9696404, -0.4002554, -0.9589243, -0.2794155, + 0.5816213, 0.96271384, 0.6569866, 0.98935825, 0.44567564, -0.37523547], + ("SymplecticAttentionP symmetric", Float64) => + [0.1411200080598672, -0.7568024953079282, -1.386695492578295, + -0.8700605238012844, -0.9589242746631385, -0.27941549819892586, + -0.29227371506864286, -0.3088672979970183, 0.6569865987187891, + 0.9893582466233818, 1.2038117178663366, 0.7047672932432907], + ("SymplecticAttentionQ arbitrary", Float32) => + Float32[-0.21044184, -0.6830057, -0.9589243, -0.2794155, -0.94552875, -0.03272158, + 0.6569866, 0.98935825, 0.9418014, 0.77742827, 0.4121185, -0.5440211], + ("SymplecticAttentionQ arbitrary", Float64) => + [-0.1882058829183166, -1.104696131437046, -0.9589242746631385, + -0.27941549819892586, 0.22227394148269097, 0.9156654275828866, + 0.6569865987187891, 0.9893582466233818, 0.6403190366622195, + 0.9701348367162665, 0.4121184852417566, -0.5440211108893698], + ("SymplecticAttentionP arbitrary", Float32) => + Float32[0.14112, -0.7568025, -0.8734648, -0.41476828, -0.9589243, -0.2794155, + 0.5487382, 1.036893, 0.6569866, 0.98935825, 0.41916823, -0.43207875], + ("SymplecticAttentionP arbitrary", Float64) => + [0.1411200080598672, -0.7568024953079282, -0.7873425751459877, + -0.5620892342512515, -0.9589242746631385, -0.27941549819892586, + 0.5260919362282062, 1.03723924510165, 0.6569865987187891, + 0.9893582466233818, 0.31665604634580913, -0.2913655625442113], + ("LinearSymplecticAttentionQ", Float32) => + Float32[0.57222337, 0.24916464, -0.9589243, -0.2794155, -1.5363214, -0.18399632, + 0.6569866, 0.98935825, 1.1174188, 0.8917234, 0.4121185, -0.5440211], + ("LinearSymplecticAttentionQ", Float64) => + [-1.1659031369145767, -1.1196573475521183, -0.9589242746631385, + -0.27941549819892586, -0.396336443224128, 0.02972159338494712, + 0.6569865987187891, 0.9893582466233818, 0.9142287872023063, + 1.5494132028555216, 0.4121184852417566, -0.5440211108893698], + ("LinearSymplecticAttentionP", Float32) => + Float32[0.14112, -0.7568025, -0.6474138, 0.88023853, -0.9589243, -0.2794155, + 1.0716593, 1.0834073, 0.6569866, 0.98935825, -0.09933251, -1.9612353], + ("LinearSymplecticAttentionP", Float64) => + [0.1411200080598672, -0.7568024953079282, -1.4240926853630695, + -0.23392131256189594, -0.9589242746631385, -0.27941549819892586, + 2.1575653701586077, 1.5290592319029808, 0.6569865987187891, + 0.9893582466233818, -0.2250659467868112, -0.8279474312485386], + ("VolumePreservingLowerLayer bias", Float32) => + Float32[0.14112, -0.6476165, -0.598324, -0.8495872, -0.9589243, -0.91151136, + 0.96390146, 1.8946021, 0.6569866, 1.4595774, -0.15381914, -1.0818557], + ("VolumePreservingLowerLayer bias", Float64) => + [0.1411200080598672, -0.7816055221167539, -1.609386788360942, + 0.5929329365287943, -0.9589242746631385, -0.11242077189834263, + 1.177798457899805, 0.2565047064153442, 0.6569865987187891, + 0.8743744480237453, 0.6992978905179187, -1.055507621880288], + ("VolumePreservingLowerLayer no bias", Float32) => + Float32[0.14112, -0.8296187, -0.5846412, -0.15975541, -0.9589243, 0.17929009, + -0.011756182, 1.5224324, 0.6569866, 0.6622396, 0.68454766, -1.0916417], + ("VolumePreservingLowerLayer no bias", Float64) => + [0.1411200080598672, -0.6244637453263984, -0.8956010555188995, + -0.08869143867827703, -0.9589242746631385, -0.9979266335643062, + 1.3230568619719785, 0.15223666690512505, 0.6569865987187891, + 1.5403054351475896, -0.21233836474850504, 0.12843598116952215], + ("VolumePreservingUpperLayer bias", Float32) => + Float32[-0.05398372, -0.14783973, -0.7908397, -0.2794155, -1.2780112, -0.6341647, + 0.11932659, 0.98935825, 1.0974605, 0.6106067, 0.7309947, -0.5440211], + ("VolumePreservingUpperLayer bias", Float64) => + [0.5246724207605964, -0.2633683979757506, -1.0024423547592334, + -0.27941549819892586, -1.432036830196382, -0.5253184460299379, + 0.8099630513669718, 0.9893582466233818, 0.6806236597388509, + 0.6693606712851976, 0.3275379323434271, -0.5440211108893698], + ("VolumePreservingUpperLayer no bias", Float32) => + Float32[0.0056689978, -1.3888329, -1.1418923, -0.2794155, -1.2376934, 0.3486507, + 1.2321696, 0.98935825, 1.015007, 1.1190019, 0.06664431, -0.5440211], + ("VolumePreservingUpperLayer no bias", Float64) => + [0.24976185047677477, -1.1778491102656683, -0.8685921488735657, + -0.27941549819892586, -0.7608309435603773, -0.49963589147857146, + 0.34682563259452204, 0.9893582466233818, 0.38764575027407044, + 1.5510517424431256, 0.5866697084870867, -0.5440211108893698], + ("VolumePreservingAttention skew_sym", Float32) => + Float32[0.8449452, 0.14761347, -0.68543327, -0.8882959, -0.7883548, -1.0432857, + -0.33902463, 0.6769342, -0.18880266, 0.72070634, 0.96760124, 0.32488802], + ("VolumePreservingAttention skew_sym", Float64) => + [0.9236902273944373, 0.46106048057839444, -0.42546614579207903, + -0.9208211598589753, -0.5445716734313131, -1.1001825596471981, + -0.6442906742752734, 0.4039590857266604, -0.4704466690511886, + 0.4546145030646437, 0.9617051976250395, 0.584608568619724], + ("VolumePreservingAttention arbitrary", Float32) => + Float32[0.8955094, -0.0821895, -0.98432374, -0.9814753, -0.13838604, 0.2740093, + 0.4344817, 0.19549364, 0.74161595, 1.2441127, 0.6027781, -0.59274805], + ("VolumePreservingAttention arbitrary", Float64) => + [1.1543908450193163, 0.6638656485499281, -0.4370145636229678, + -1.136105601396825, -0.1849747067064138, 0.13902306355623095, + 0.33520367032298276, 0.22319956846571176, 0.06517781005519652, + 1.0814840572909876, 1.1034788497727075, 0.11094027672684503], + ("MultiHeadAttention 1 head", Float32) => + Float32[-0.2477576, -0.5865091, -1.0242593, -0.5696287, -0.8753681, -0.29568246, + 0.7043326, 0.92572075, 0.8710045, 0.8627129, 0.3941149, -0.1801846], + ("MultiHeadAttention 1 head", Float64) => + [0.24130465306760224, -0.968280864456187, -1.1109616621891816, + -0.15876018724605234, -0.6382017204253616, -0.22947608567514652, + 0.5119821347757704, 1.0567444358549594, 0.5381358914333416, + 1.3082837856985623, 0.624468846621764, -0.7170039529551228], + ("MultiHeadAttention 2 heads", Float32) => + Float32[0.06764289, -0.6196884, -1.0301708, -0.14345966, -0.9394911, -0.31636807, + 0.68966633, 0.68806624, 0.6494659, 1.006897, 0.33891702, -0.4590907], + ("MultiHeadAttention 2 heads", Float64) => + [0.1825661813180538, -0.24253428457883852, -0.9594250287721061, + -0.26116560971643793, -0.999233977123587, -0.4760490212094931, + -0.014988601998192408, 0.9637506951879614, 0.6696001107858108, + 0.9950399041752966, 0.7647174325787416, -0.47275268235644763], + ("MultiHeadAttention 1 head Stiefel", Float32) => + Float32[-0.08497381, -1.0349782, -0.51467085, -0.48130214, -0.8670119, -0.13992026, + 0.30621412, 1.2289094, 0.89886147, 1.229189, 0.3090373, -0.67246187], + ("MultiHeadAttention 1 head Stiefel", Float64) => + [-0.08497377552471627, -1.0349781542295124, -0.5146708127126829, + -0.4813021790408115, -0.867011877934448, -0.1399203040897155, + 0.3062141834726896, 1.2289093420989798, 0.8988614583360413, + 1.2291890122715268, 0.30903727120923385, -0.6724618412705963], + ("MultiHeadAttention 2 heads Stiefel", Float32) => + Float32[0.07385998, -0.7862042, -0.9583907, -0.36726886, -0.8535721, -0.25491065, + 0.9084542, 0.89053684, 0.5968184, 0.86706436, 0.106292814, -0.16230369], + ("MultiHeadAttention 2 heads Stiefel", Float64) => + [0.07385999041019796, -0.7862041928808658, -0.9583906866557481, + -0.3672688414136106, -0.8535720566659938, -0.2549106752878051, + 0.9084541348511641, 0.8905368659247743, 0.5968183653182471, + 0.8670643266791885, 0.10629277175877594, -0.16230365424515353], + ("ResNetLayer bias", Float32) => + Float32[0.6245465, -1.0234065, -1.813796, -1.0477115, -0.8991417, 0.01658684, + 1.5749986, 0.78809625, 0.13631392, 1.0086006, 0.37422135, 0.2853077], + ("ResNetLayer bias", Float64) => + [0.7071374945546426, -1.4985667084062746, -0.7522490077073039, + -0.24820819080854345, -1.8552453015305264, -0.5637852544592185, + 0.13546964359788516, 0.2629915262503477, 1.170753359274967, + 1.8223338173748993, 0.6773435457894211, 0.08228756562282324], + ("ResNetLayer no bias", Float32) => + Float32[-0.102672964, -0.56263745, -1.2477826, 0.34026194, -1.5591141, 0.59330374, + 1.1851288, 0.44095588, 1.3352829, 0.123678505, 0.22272642, -0.75262415], + ("ResNetLayer no bias", Float64) => + [-0.10920357251312326, -0.09477426958474733, -1.5525068731619602, + -0.5521662467036292, -0.6952562579390011, -1.0804326347863957, + 1.2793706873835915, 1.059427508994182, 0.6879767504521507, + 1.1090902733818453, 0.48850082542885215, -0.3261506753230139], + ("PSDLayer down", Float32) => + Float32[-0.7568025, -0.2794155, -0.2794155, 0.98935825, 0.98935825, -0.5440211], + ("PSDLayer down", Float64) => + [-0.7568024953079282, -0.27941549819892586, -0.27941549819892586, + 0.9893582466233818, 0.9893582466233818, -0.5440211108893698], + ("PSDLayer up", Float32) => Float32[0.0, 0.14112, 0.0, -0.7568025, 0.0, -0.9589243, + 0.0, -0.2794155, 0.0, 0.6569866, 0.0, 0.98935825], + ("PSDLayer up", Float64) => + [0.0, 0.1411200080598672, 0.0, -0.7568024953079282, 0.0, -0.9589242746631385, + 0.0, -0.27941549819892586, 0.0, 0.6569865987187891, 0.0, 0.9893582466233818], + ("PositionalEncoding", Float32) => + Float32[0.14112, 0.2431975, -0.9589243, 0.7205845, -0.11745334, 0.2608868, + 0.6669864, 1.9893082, 1.566284, 0.57321143, 0.43211716, 0.4557789], + ("PositionalEncoding", Float64) => + [0.1411200080598672, 0.2431975046920718, -0.9589242746631385, + 0.7205845018010741, -0.11745328985524195, 0.2608868076692139, + 0.6669864320529557, 1.989308247040047, 1.5662840255444708, + 0.5732114100762393, 0.4321171519350897, 0.455778895777208] +) + +const ARCHITECTURE_PINS = Dict{Tuple{String, DataType}, Vector}( + ("GSympNet upper", Float32) => + Float32[0.45778942, -0.47966588, -1.084514, -0.1737431, -1.1854849, -0.82048124, + 0.6996686, 0.898326, 0.6795032, 1.1800716, 0.49674934, -0.59732735], + ("GSympNet upper", Float64) => + [0.7812255409893067, 0.0173785773736147, -1.5171452421313472, + -0.019845341221890023, -1.7232505009858454, -1.192378913862119, + 1.192294274750766, 0.8894430266076818, 0.8145733014779291, + 1.1486632113960913, 0.2527785182370254, -0.7614967882863601], + ("GSympNet lower, 3 layers", Float32) => + Float32[0.5086076, -0.47647104, -0.8954462, -0.34905446, -1.5385884, -0.33480173, + 1.1682279, 1.4478953, 0.80730814, 0.7930634, -0.119695544, -0.9036403], + ("GSympNet lower, 3 layers", Float64) => + [0.15514484976305312, -0.4569952608455856, -0.8611083634264115, + -0.1056093039823252, -1.2057557150300404, -1.5480309422568608, + 1.0206388488571727, 1.1254905381332188, 0.8765345002088953, + 1.919845435840653, 0.05472598607052265, -0.8886720815032787], + ("LinearSymplecticTransformer upper", Float32) => + Float32[0.31126082, -1.0972345, -1.5355496, -0.98436165, -1.1455917, 0.32650018, + -0.8194469, 2.235367, -0.6862323, 0.82518685, 0.9095292, 0.15315211], + ("LinearSymplecticTransformer upper", Float64) => + [0.3682044343180707, -1.5718809942889087, -2.891573283590699, + -1.2885739507034084, -1.943728776995327, -1.3818309566338627, + 1.8827824360961185, 1.7997944349933976, 0.15179423847796197, + 0.7856506516411316, -1.6050282052100522, -2.5522896032992644], + ("LinearSymplecticTransformer lower", Float32) => + Float32[-0.39304054, -0.281504, -0.63551927, 0.8481669, -1.1012809, 0.114686936, + 1.2375215, 0.7951687, 1.4403617, 1.0673912, -2.250509, -2.599523], + ("LinearSymplecticTransformer lower", Float64) => + [-0.5574855538571216, -1.7042882025431472, -2.2697583569120248, + -0.3249139343198121, -0.29920846559061787, 0.629611530717385, + 1.0962623221293337, 1.1629909624038701, 1.0207042927011911, + -0.408155304285686, 1.294964910373253, 1.4057308875592829], + ("SymplecticTransformer", Float32) => + Float32[-0.6765071, -2.0309815, 0.5501745, -0.16158332, -0.39364353, 1.300374, + -1.5025885, 2.506516, 0.59647775, 0.38502067, 0.9785921, -1.2965591], + ("SymplecticTransformer", Float64) => + [-1.6696113747905288, -2.8385769261213123, 0.8881083499247768, + -2.1144242957185004, 0.10910065819940973, 1.5097746566910166, + -2.027167467302565, 3.544763300787557, 0.5567870150274278, + 0.6900606391002133, -0.1336398759069721, -0.36446155956788623], + ("SymplecticTransformer upscaled, arbitrary, lower", Float32) => + Float32[-0.3829862, 0.010057211, -0.2560426, -0.46741015, -0.6697832, 0.2807653, + 0.39536962, 0.23018713, 0.9805062, -0.6856048, -0.33859593, 0.5296659], + ("SymplecticTransformer upscaled, arbitrary, lower", Float64) => + [-0.14172800925450854, 0.47366485362175237, -0.4361261342070305, + 0.37191956660278486, -0.40197966711195865, 0.363051103749285, + 0.7427700227572049, 0.31112315194353124, 0.527994091100511, + -0.638322741786857, -0.31384758827811815, -0.5039319672266002], + ("StandardTransformerIntegrator 2 heads", Float32) => + Float32[0.022821024, -0.69257545, -2.1877477, 0.91556484, -0.52159727, 0.6637138, + 1.5511532, 0.15770149, 0.14093326, 0.6675812, 0.028205752, 0.6148718], + ("StandardTransformerIntegrator 2 heads", Float64) => + [0.17977908729262387, 0.12169866081646796, -1.6231521596339584, + 0.7857606141059303, 4.670034697539226, -2.69094043368872, + -10.112809523830482, 5.906810200145711, -0.45339263819616205, + 0.18518106200648266, 4.387579005868016, -2.902958442828254], + ("StandardTransformerIntegrator upscaled", Float32) => + Float32[-1.8204737, 0.022696197, 0.32511562, 1.0635418, -3.2326684, -1.3320627, + 1.7759933, 2.476647, -3.2956362, -1.3550298, 1.2212591, 2.8090253], + ("StandardTransformerIntegrator upscaled", Float64) => + [0.5608860920764921, 1.869890865597026, 1.5275026424666613, + 6.586190962790901, -1.3720026027505405, 0.3051354144292979, + -2.3455660108879925, -0.10747437011986571, -2.1095145343690547, + 3.7843161588112997, -0.952158125884726, 1.5066098615395283], + ("VolumePreservingFeedForward", Float32) => + Float32[-0.72681475, -0.88861287, 0.35451126, 0.88032955, 0.9107045, 0.69980294, + 0.44388264, -0.863206, -0.41037926, 0.55101925, -1.0566202, -0.102654755], + ("VolumePreservingFeedForward", Float64) => + [-0.09568708623180941, 0.033142006733737406, 1.5201120238992119, + 1.1536442884174085, -1.4572121348990592, 0.010690993799266635, + 2.12100273651203, 1.227044050748543, 1.1691130377482357, + -0.035159691827130635, -3.326019464710985, -2.1485119360160394], + ("ResNet", Float32) => + Float32[0.73120904, 0.8714113, -3.8866358, 0.06171602, -2.7615, -4.03835, + 3.8442721, -1.8685002, 1.6793127, 1.9793546, -0.41924566, 1.4315214], + ("ResNet", Float64) => [-0.646345971756664, -0.23753121190813054, -0.9951787204212194, + -0.2704252267795757, -0.1671496861980032, -0.03688570424392568, + 0.7204417548764077, -0.5221018094099064, 0.5394184873010252, + 0.29677436280536323, 0.3275755158099256, 0.5605760656369991], + ("VolumePreservingTransformer", Float32) => + Float32[-3.6292968, -3.405287, -1.1849939, -1.8699291, 2.5299454, 0.2920828, + 3.3345196, 1.9946433, 0.582687, -0.07292181, 0.9494109, 0.46535522], + ("VolumePreservingTransformer", Float64) => + [-0.5976912408192199, -1.3144611021512498, 0.3929764864245985, + -1.5663516954047316, 0.7612461396190153, 1.5861817535176457, + 0.8475345446546175, 0.4037465552433255, -0.39355343624995776, + -0.9254715244419354, 0.9227445284119706, -1.9991973742091549], + ("SymplecticAutoencoder", Float32) => + Float32[0.14824425, 1.4241984, 0.0, 1.2083642, 0.0, -1.1983773, 0.0, -0.5167034, + -0.11093118, -0.24777675, 0.0, -0.6997613, 0.0, 0.93032384, 0.0, 0.8809319, + -0.062251065, -1.289555, 0.0, -0.6230452, 0.0, 0.43147907, 0.0, -0.18738672], + ("SymplecticAutoencoder", Float64) => + [0.12425930698294496, 0.6236570498812481, 0.0, 1.2761302986710867, 0.0, + -0.8215585135050352, 0.0, -0.16738782714668204, -0.2164377448101824, + 0.15312826507468996, 0.0, -1.0588365432658748, 0.0, 1.0526587337560476, + 0.0, 0.8404268414026151, 0.08015720766406666, -0.7301872206984017, 0.0, + -0.29388994566905396, 0.0, -0.06862981365887028, 0.0, -0.4969405960745537], + ("PSDArch", Float32) => + Float32[0.0, 0.0, 0.7568025, 0.9589243, 0.0, 0.0, -0.98935825, -0.4121185, + 0.0, 0.0, 0.2794155, -0.6569866, 0.0, 0.0, 0.5440211, 0.9999902, 0.0, + 0.0, -0.98935825, -0.4121185, 0.0, 0.0, 0.53657293, -0.42016703], + ("PSDArch", Float64) => + [0.0, 0.0, 0.7568024953079282, 0.9589242746631385, 0.0, 0.0, -0.9893582466233818, + -0.4121184852417566, 0.0, 0.0, 0.27941549819892586, -0.6569865987187891, 0.0, + 0.0, 0.5440211108893698, 0.9999902065507035, 0.0, 0.0, -0.9893582466233818, + -0.4121184852417566, 0.0, 0.0, 0.5365729180004349, -0.4201670368266409] +) + +const AUTOENCODER_PINS = Dict{Tuple{String, Symbol, DataType}, Vector}( + ("SymplecticAutoencoder", :encoder, Float32) => + Float32[-0.4171233, 1.0425391, 0.5292804, -0.17343941, -0.21139398, -0.38436216, + -0.5211674, 0.9726173, 0.64541554, -0.672654, -0.087599784, -0.64801896], + ("SymplecticAutoencoder", :decoder, Float32) => + Float32[0.035465937, -0.19175482, 0.0, -0.3003301, 0.0, 0.9589243, 0.0, -0.2794155, + -0.20066342, 1.324945, 0.0, -0.70665956, 0.0, -0.6569866, 0.0, 0.98935825, + 0.13391133, -0.89653385, 0.0, 0.92782503, 0.0, -0.4121185, 0.0, -0.5440211], + ("SymplecticAutoencoder", :encoder, Float64) => + [-0.38456057024632395, 0.8887900593334842, 0.6862749160321735, + 0.36313989724509343, -0.5534771654378703, -0.6524480692813724, + -0.43712546050074197, 0.21780980894729396, 0.820779383759808, + -0.3603089380106658, -0.2603333446594144, -0.5310485388343394], + ("SymplecticAutoencoder", :decoder, Float64) => + [-0.17883285680270908, -0.29674852059978074, 0.0, -0.8197623649599224, + 0.0, 0.9589242746631385, 0.0, -0.27941549819892586, 0.0499168635334002, + 0.982124348823736, 0.0, -0.35743997859239146, 0.0, -0.6569865987187891, + 0.0, 0.9893582466233818, 0.14549272300750837, -0.5166795220675615, 0.0, + 1.1154992226593332, 0.0, -0.4121184852417566, 0.0, -0.5440211108893698], + ("PSDArch", :encoder, Float32) => + Float32[-0.7568025, 0.9589243, 0.98935825, -0.4121185, -0.2794155, -0.6569866, + -0.5440211, 0.9999902, 0.98935825, -0.4121185, -0.53657293, -0.42016703], + ("PSDArch", :decoder, Float32) => + Float32[0.0, 0.0, -0.14112, -0.7568025, 0.0, 0.0, 0.9589243, -0.2794155, + 0.0, 0.0, 0.9589243, -0.2794155, 0.0, 0.0, -0.6569866, 0.98935825, + 0.0, 0.0, -0.6569866, 0.98935825, 0.0, 0.0, -0.4121185, -0.5440211], + ("PSDArch", :encoder, Float64) => + [-0.7568024953079282, 0.9589242746631385, 0.9893582466233818, + -0.4121184852417566, -0.27941549819892586, -0.6569865987187891, + -0.5440211108893698, 0.9999902065507035, 0.9893582466233818, + -0.4121184852417566, -0.5365729180004349, -0.4201670368266409], + ("PSDArch", :decoder, Float64) => + [0.0, 0.0, -0.1411200080598672, -0.7568024953079282, 0.0, 0.0, + 0.9589242746631385, -0.27941549819892586, 0.0, 0.0, 0.9589242746631385, + -0.27941549819892586, 0.0, 0.0, -0.6569865987187891, + 0.9893582466233818, 0.0, 0.0, -0.6569865987187891, 0.9893582466233818, + 0.0, 0.0, -0.4121184852417566, -0.5440211108893698] +) + +const LASYMPNET_LAYER_TYPES = Dict{Tuple{Bool, Bool}, Vector{String}}( + (true, true) => + ["LinearLayer{4, 4, :Q}", "LinearLayer{4, 4, :P}", "LinearLayer{4, 4, :Q}", + "BiasLayer{4, 4}", "ActivationLayer{4, 4, typeof(tanh), :Q}", + "ActivationLayer{4, 4, typeof(tanh), :P}", "LinearLayer{4, 4, :Q}", + "LinearLayer{4, 4, :P}", "LinearLayer{4, 4, :Q}", + "BiasLayer{4, 4}", "ActivationLayer{4, 4, typeof(tanh), :Q}", + "ActivationLayer{4, 4, typeof(tanh), :P}", "LinearLayer{4, 4, :Q}", + "LinearLayer{4, 4, :P}", "LinearLayer{4, 4, :Q}", "BiasLayer{4, 4}"], + (true, false) => + ["LinearLayer{4, 4, :Q}", "LinearLayer{4, 4, :P}", "LinearLayer{4, 4, :Q}", + "BiasLayer{4, 4}", "ActivationLayer{4, 4, typeof(tanh), :P}", + "ActivationLayer{4, 4, typeof(tanh), :Q}", "LinearLayer{4, 4, :Q}", + "LinearLayer{4, 4, :P}", "LinearLayer{4, 4, :Q}", + "BiasLayer{4, 4}", "ActivationLayer{4, 4, typeof(tanh), :P}", + "ActivationLayer{4, 4, typeof(tanh), :Q}", "LinearLayer{4, 4, :Q}", + "LinearLayer{4, 4, :P}", "LinearLayer{4, 4, :Q}", "BiasLayer{4, 4}"], + (false, true) => + ["LinearLayer{4, 4, :P}", "LinearLayer{4, 4, :Q}", "LinearLayer{4, 4, :P}", + "BiasLayer{4, 4}", "ActivationLayer{4, 4, typeof(tanh), :Q}", + "ActivationLayer{4, 4, typeof(tanh), :P}", "LinearLayer{4, 4, :P}", + "LinearLayer{4, 4, :Q}", "LinearLayer{4, 4, :P}", + "BiasLayer{4, 4}", "ActivationLayer{4, 4, typeof(tanh), :Q}", + "ActivationLayer{4, 4, typeof(tanh), :P}", "LinearLayer{4, 4, :P}", + "LinearLayer{4, 4, :Q}", "LinearLayer{4, 4, :P}", "BiasLayer{4, 4}"], + (false, false) => + ["LinearLayer{4, 4, :P}", "LinearLayer{4, 4, :Q}", "LinearLayer{4, 4, :P}", + "BiasLayer{4, 4}", "ActivationLayer{4, 4, typeof(tanh), :P}", + "ActivationLayer{4, 4, typeof(tanh), :Q}", "LinearLayer{4, 4, :P}", + "LinearLayer{4, 4, :Q}", "LinearLayer{4, 4, :P}", + "BiasLayer{4, 4}", "ActivationLayer{4, 4, typeof(tanh), :P}", + "ActivationLayer{4, 4, typeof(tanh), :Q}", "LinearLayer{4, 4, :P}", + "LinearLayer{4, 4, :Q}", "LinearLayer{4, 4, :P}", "BiasLayer{4, 4}"] +) diff --git a/test/helpers/pins.jl b/test/helpers/pins.jl new file mode 100644 index 000000000..f96249b17 --- /dev/null +++ b/test/helpers/pins.jl @@ -0,0 +1,150 @@ +# The cases of the pinned forward passes, and the comparison that a pin uses. +# +# A pin is the output of one forward pass of a seeded network on a fixed input, recorded as a +# literal from commit 4c4a8320 with the cases and the seeds below. The literals are in +# `test/helpers/pinned_outputs.jl`. +# +# Two things keep a pin independent of the platform: +# - the input is a closed form, `sin(i + 2j)`, and not a random draw; +# - every `StiefelManifold` leaf is overwritten with a fixed signed selection of columns of the +# identity before the pass. Its seeded initialisation is orthonormalised through BLAS, whose last +# bits differ between OpenBLAS versions. +# The comparison is at 16 ULPs of the largest magnitude of the pinned vector. A value that passes +# through GEMM and `tanh` differs by a few ULPs between x86_64 and arm64, and a wrong layer moves it +# by a fraction. The scale is the vector's and not each entry's: an entry that comes out of +# cancellation, such as -0.0686 beside 1.276 in the deep `SymplecticAutoencoder` pass, carries the +# rounding error of the terms it came from, and floating point gives it no accuracy relative to its +# own magnitude. + +using GeometricMachineLearning +using GeometricMachineLearning: params, ResNetLayer +using NeuralNetworkParameters: NetworkParameters +import Random + +const PIN_ULPS = 16 + +function pin_input(::Type{T}, dim::Integer, cols::Integer) where {T} + T[sin(i + 2j) + for i in 1:dim, j in 1:cols] +end + +# The deviation of `a` from the pin `b`, in ULPs of the largest magnitude of `b`: the largest +# `abs(x - y)` over the entries, divided by `eps(maximum(abs, b))`. `Inf` when the sizes differ. +function pin_ulps(a::AbstractArray, b::AbstractArray) + size(a) == size(b) || return Inf + maximum(abs(x - y) for (x, y) in zip(a, b)) / eps(maximum(abs, b)) +end + +# A type as its name and its parameters, so that a pin of a type does not depend on the alias that +# `show` picks for it. +type_string(t::DataType) = string(nameof(t), "{", join(repr.(t.parameters), ", "), "}") + +# Overwrite `A` with the signed columns `±e_{j + shift}` of the identity, so that `A` has exactly +# orthonormal columns and a different pattern for each `shift`. +function fixed_orthonormal!(A::AbstractMatrix, shift::Integer) + n, k = size(A) + @assert k ≤ n + fill!(A, zero(eltype(A))) + for j in 1:k + A[mod1(j + shift, n), j] = iseven(j + shift) ? one(eltype(A)) : -one(eltype(A)) + end + A +end + +function _fix_stiefel!(x::StiefelManifold, count::Base.RefValue{Int}) + fixed_orthonormal!( + x.A, count[] += 1) +end +function _fix_stiefel!(x::Union{NamedTuple, NetworkParameters}, count::Base.RefValue{Int}) + for k in keys(x) + _fix_stiefel!(x[k], count) + end +end +_fix_stiefel!(::Any, ::Base.RefValue{Int}) = nothing + +# Overwrite every Stiefel leaf of `ps`, each with its own pattern, in the order of the keys. +fix_stiefel_leaves!(ps) = (_fix_stiefel!(ps, Ref(0)); ps) + +# A seeded network of element type `T` for `model`, with its Stiefel leaves fixed. +function pinned_network(model, ::Type{T}, seed::Integer) where {T} + Random.seed!(seed) + nn = NeuralNetwork(model, CPU(), T) + fix_stiefel_leaves!(params(nn)) + nn +end + +# The seed of the `i`-th case of each list below. +layer_seed(i::Integer) = 1000 + i +architecture_seed(i::Integer) = 2000 + i +autoencoder_seed(i::Integer) = 3000 + i + +# Each case is `(name, model, input dimension, number of columns)`. A layer is wrapped in a `Chain`. +function layer_pin_cases() + (("GradientLayerQ", Chain(GradientLayerQ(4, 6, tanh)), 4, 3), + ("GradientLayerP", Chain(GradientLayerP(4, 6, tanh)), 4, 3), + ("ActivationLayerQ", Chain(ActivationLayerQ(4, tanh)), 4, 3), + ("ActivationLayerP", Chain(ActivationLayerP(4, tanh)), 4, 3), + ("LinearLayerQ", Chain(LinearLayerQ(4)), 4, 3), + ("LinearLayerP", Chain(LinearLayerP(4)), 4, 3), + ("SymplecticAttentionQ symmetric", + Chain(SymplecticAttentionQ(4; symmetric = true)), 4, 3), + ("SymplecticAttentionP symmetric", + Chain(SymplecticAttentionP(4; symmetric = true)), 4, 3), + ("SymplecticAttentionQ arbitrary", + Chain(SymplecticAttentionQ(4; symmetric = false)), 4, 3), + ("SymplecticAttentionP arbitrary", + Chain(SymplecticAttentionP(4; symmetric = false)), 4, 3), + ("LinearSymplecticAttentionQ", Chain(LinearSymplecticAttentionQ(4, 3)), 4, 3), + ("LinearSymplecticAttentionP", Chain(LinearSymplecticAttentionP(4, 3)), 4, 3), + ("VolumePreservingLowerLayer bias", + Chain(VolumePreservingLowerLayer(4, tanh; use_bias = true)), 4, 3), + ("VolumePreservingLowerLayer no bias", + Chain(VolumePreservingLowerLayer(4, tanh; use_bias = false)), 4, 3), + ("VolumePreservingUpperLayer bias", + Chain(VolumePreservingUpperLayer(4, tanh; use_bias = true)), 4, 3), + ("VolumePreservingUpperLayer no bias", + Chain(VolumePreservingUpperLayer(4, tanh; use_bias = false)), 4, 3), + ("VolumePreservingAttention skew_sym", + Chain(VolumePreservingAttention(4, 3; skew_sym = true)), 4, 3), + ("VolumePreservingAttention arbitrary", + Chain(VolumePreservingAttention(4, 3; skew_sym = false)), 4, 3), + ("MultiHeadAttention 1 head", Chain(MultiHeadAttention(4, 1)), 4, 3), + ("MultiHeadAttention 2 heads", Chain(MultiHeadAttention(4, 2)), 4, 3), + ("MultiHeadAttention 1 head Stiefel", + Chain(MultiHeadAttention(4, 1; Stiefel = true)), 4, 3), + ("MultiHeadAttention 2 heads Stiefel", + Chain(MultiHeadAttention(4, 2; Stiefel = true)), 4, 3), + ("ResNetLayer bias", Chain(ResNetLayer(4, tanh; use_bias = true)), 4, 3), + ("ResNetLayer no bias", Chain(ResNetLayer(4, tanh; use_bias = false)), 4, 3), + ("PSDLayer down", Chain(PSDLayer(4, 2)), 4, 3), + ("PSDLayer up", Chain(PSDLayer(2, 4)), 2, 3), + ("PositionalEncoding", Chain(PositionalEncoding(4)), 4, 3)) +end + +function architecture_pin_cases() + (("GSympNet upper", GSympNet(4), 4, 3), + ("GSympNet lower, 3 layers", GSympNet(4; n_layers = 3, init_upper = false), 4, 3), + ("LinearSymplecticTransformer upper", LinearSymplecticTransformer(4, 3), 4, 3), + ("LinearSymplecticTransformer lower", + LinearSymplecticTransformer(4, 3; init_upper = false, n_sympnet = 3), 4, 3), + ("SymplecticTransformer", SymplecticTransformer(4), 4, 3), + ("SymplecticTransformer upscaled, arbitrary, lower", + SymplecticTransformer(4; transformer_dim = 6, symmetric = false, init_upper = false), + 4, 3), + ("StandardTransformerIntegrator 2 heads", + StandardTransformerIntegrator(4; n_heads = 2), 4, 3), + ("StandardTransformerIntegrator upscaled", + StandardTransformerIntegrator(4; transformer_dim = 6, n_heads = 2), 4, 3), + ("VolumePreservingFeedForward", VolumePreservingFeedForward(4), 4, 3), + ("ResNet", ResNet(4, 2, tanh), 4, 3), + ("VolumePreservingTransformer", VolumePreservingTransformer(4, 3), 4, 3), + ("SymplecticAutoencoder", SymplecticAutoencoder(8, 4), 8, 3), + ("PSDArch", PSDArch(8, 4), 8, 3)) +end + +# Each case is `(name, architecture, full dimension, reduced dimension)`; the pins are the outputs of +# `encoder(nn)` on the full input and of `decoder(nn)` on the reduced input. +function autoencoder_pin_cases() + (("SymplecticAutoencoder", SymplecticAutoencoder(8, 4), 8, 4), + ("PSDArch", PSDArch(8, 4), 8, 4)) +end diff --git a/test/integration/documented_defaults.jl b/test/integration/documented_defaults.jl index 47880b7a4..73ddbc751 100644 --- a/test/integration/documented_defaults.jl +++ b/test/integration/documented_defaults.jl @@ -40,7 +40,7 @@ end @test arch.n_blocks === GML.sti_n_blocks_default @test arch.L === GML.sti_L_default - @test arch.upsacling_activation === GML.sti_upscaling_activation_default + @test arch.upscaling_activation === GML.sti_upscaling_activation_default @test arch.resnet_activation === GML.sti_resnet_activation_default @test arch.attention_activation === GML.sti_attention_activation_default @test arch.add_connection === GML.sti_add_connection_default diff --git a/test/layers/layer_pins.jl b/test/layers/layer_pins.jl new file mode 100644 index 000000000..214f8d5ba --- /dev/null +++ b/test/layers/layer_pins.jl @@ -0,0 +1,31 @@ +# The forward pass of each layer whose methods share code across a `Q`/`P` pair, a `Bool` choice or +# a triangle, against the output that commit 4c4a8320 records for the same seed and input. The +# cases cover dense, `SymmetricMatrix`, triangular, `SkewSymMatrix` and Stiefel leaves; the helper +# file says how the pins are kept independent of the platform. + +using GeometricMachineLearning +using Test + +include("../helpers/eltypes.jl") +include("../helpers/pins.jl") +include("../helpers/pinned_outputs.jl") + +@testset "$name ($T)" for (i, (name, model, dim, cols)) in enumerate(layer_pin_cases()), + T in REAL_ELTYPES + + nn = pinned_network(model, T, layer_seed(i)) + y = nn(pin_input(T, dim, cols)) + @test eltype(y) === T + @test pin_ulps(vec(y), LAYER_PINS[(name, T)]) ≤ PIN_ULPS +end + +# The comparison itself: its scale is the largest magnitude of the pin, not the entry's own, and a +# deviation of 16 of those ULPs passes where 17 fail. +@testset "pin_ulps counts in ULPs of the largest magnitude ($T)" for T in REAL_ELTYPES + x = T[1.5, -0.25, 3] + shifted(n) = [x[1], x[2] + n * eps(T(3)), x[3]] + @test pin_ulps(shifted(16), x) == 16 + @test pin_ulps(shifted(16), x) ≤ PIN_ULPS + @test !(pin_ulps(shifted(17), x) ≤ PIN_ULPS) + @test pin_ulps(x[1:2], x) == Inf +end diff --git a/test/layers/layer_structure.jl b/test/layers/layer_structure.jl new file mode 100644 index 000000000..38ee37ed1 --- /dev/null +++ b/test/layers/layer_structure.jl @@ -0,0 +1,251 @@ +# The volume-preserving, symplectic-attention and ResNet layers and the symplectic autoencoder +# encode each structural choice as a `Symbol` type parameter, chosen in one expression from the +# public `Bool` keyword; a `Q`/`P` pair, a lower/upper pair and the two weight kinds of +# multi-head attention share one method; the SympNet layers store their dimensions as `Int`; and +# every type argument under `src/layers/` is spelled `::Type{T}`. + +using GeometricMachineLearning +using GeometricMachineLearning: GradientLayer, LinearLayer, ResNetLayer, initialparameters, + params +using AbstractNeuralNetworks: parameterlength, ZeroInitializer, OneInitializer +using LinearAlgebra: I +using Test +import Random + +const GML = GeometricMachineLearning + +include("../helpers/eltypes.jl") + +# The methods of `f` whose definitions are in the file `file` under `src/`. +function methods_in(f, file::AbstractString) + filter( + m -> endswith(String(m.file), joinpath("src", file)), collect(methods(f))) +end + +# The forward methods on a `(q, p)` input that the layers `layers` reach in the file `file`: the +# methods whose input argument is not an array. +function qp_forward_methods(layers, file::AbstractString) + found = Set{Method}() + for l in layers, m in methods_in(l, file) + + input = Base.unwrap_unionall(m.sig).parameters[2] + input <: AbstractArray || push!(found, m) + end + found +end + +@testset "a Q/P pair shares its forward method" begin + sympnet_layers = (GradientLayerQ(4, 6, tanh), GradientLayerP(4, 6, tanh), + ActivationLayerQ(4, tanh), ActivationLayerP(4, tanh), LinearLayerQ(4), LinearLayerP(4)) + @test length(qp_forward_methods(sympnet_layers, joinpath("layers", "sympnets.jl"))) == 3 + + attention_layers = (SymplecticAttentionQ(4; symmetric = true), + SymplecticAttentionP(4; symmetric = true), SymplecticAttentionQ(4; symmetric = false), + SymplecticAttentionP(4; symmetric = false)) + @test length(qp_forward_methods(attention_layers, + joinpath("layers", "symplectic_attention.jl"))) == 2 + + linear_attention_layers = ( + LinearSymplecticAttentionQ(4, 3), LinearSymplecticAttentionP(4, 3)) + @test length(qp_forward_methods(linear_attention_layers, + joinpath("layers", "linear_symplectic_attention.jl"))) == 1 +end + +# Julia drops a type variable that the body of an alias does not use, so the alias type is the same +# with or without it, and an unused variable shows only in the source text. +@testset "the linear SympNet aliases name two parameters" begin + source = read(joinpath(pkgdir(GML), "src", "layers", "sympnets.jl"), String) + @test occursin("const LinearLayerQ{M, N} = LinearLayer{M, N, :Q}", source) + @test occursin("const LinearLayerP{M, N} = LinearLayer{M, N, :P}", source) + @test LinearLayerQ{4, 4} === LinearLayer{4, 4, :Q} + @test LinearLayerP{4, 4} === LinearLayer{4, 4, :P} +end + +@testset "a SympNet layer raises on a half of the wrong size ($T)" for T in REAL_ELTYPES + # In a layer of dimension 2 the scale has length 1 and broadcasts against a half of any length, + # so there a half of the wrong size can pass the update unless the layer checks the dimension. + for (M, layers) in ( + (4, + (GradientLayerQ(4, 6, tanh), + GradientLayerP(4, 6, tanh), ActivationLayerQ(4, tanh), + ActivationLayerP(4, tanh), LinearLayerQ(4), LinearLayerP(4))), + (2, (ActivationLayerQ(2, tanh), ActivationLayerP(2, tanh)))) + for l in layers + ps = params(NeuralNetwork(Chain(l), CPU(), T))[1] + m = M ÷ 2 + for n in (m - 1, m + 1, m + 3) + n == 0 && continue + @test_throws ErrorException("Dimension mismatch.") l( + (q = ones(T, n, 3), p = ones(T, n, 3)), ps) + @test_throws ErrorException("Dimension mismatch.") l( + (q = ones(T, n), p = ones(T, n)), ps) + @test_throws ErrorException("Dimension mismatch.") l( + (q = ones(T, n, 3), p = ones(T, m, 3)), ps) + @test_throws ErrorException("Dimension mismatch.") l( + (q = ones(T, m, 3), p = ones(T, n, 3)), ps) + @test_throws ErrorException("Dimension mismatch.") l( + (q = ones(T, m), p = ones(T, n)), ps) + end + @test eltype(l((q = ones(T, m, 3), p = ones(T, m, 3)), ps).q) === T + end + end +end + +@testset "a SympNet layer stores its dimensions as Int" begin + for l in (GradientLayerQ(4, 8, tanh), GradientLayerQ(Int32(4), 8, tanh), + GradientLayerP(4, Int32(8), tanh)) + @test isconcretetype(fieldtype(typeof(l), :second_dim)) + @test l.second_dim === 8 + end + @test GradientLayerQ(Int32(4), Int32(8), tanh) === GradientLayerQ(4, 8, tanh) + @test GradientLayerP(Int32(4), Int32(8), tanh) === GradientLayerP(4, 8, tanh) + @test ActivationLayerQ(Int32(4), tanh) === ActivationLayerQ(4, tanh) + @test ActivationLayerP(Int32(4), tanh) === ActivationLayerP(4, tanh) + @test LinearLayerQ(Int32(4)) === LinearLayerQ(4) + @test LinearLayerP(Int32(4)) === LinearLayerP(4) +end + +@testset "each Bool keyword of a ladder becomes a Symbol type parameter" begin + for skew_sym in (true, false) + @test typeof(VolumePreservingAttention(4, 3; skew_sym)) === + VolumePreservingAttention{4, 4, skew_sym ? :skew_sym : :arbitrary, 3} + end + for symmetric in (true, false), + (constructor, side) in ((SymplecticAttentionQ, :Q), (SymplecticAttentionP, :P)) + + l = constructor(4; symmetric) + @test typeof(l).parameters[3] === side + @test typeof(l).parameters[4] === (symmetric ? :symmetric : :arbitrary) + end + for use_bias in (true, false) + bias = use_bias ? :bias : :no_bias + @test typeof(VolumePreservingLowerLayer(3; use_bias)) === + VolumePreservingLowerLayer{3, 3, bias, typeof(tanh)} + @test typeof(VolumePreservingUpperLayer(3; use_bias)) === + VolumePreservingUpperLayer{3, 3, bias, typeof(tanh)} + # `ResNetLayer` encodes the bias as the volume-preserving layers do. + @test typeof(ResNetLayer(3; use_bias)).parameters[3] === + typeof(VolumePreservingLowerLayer(3; use_bias)).parameters[3] === bias + end + @test ResNetLayer(3) isa ResNetLayer{3, 3, :bias, typeof(identity)} + for encoder_init_q in (true, false), decoder_init_q in (true, false) + + @test typeof(SymplecticAutoencoder(8, 4; encoder_init_q, decoder_init_q)) === + SymplecticAutoencoder{encoder_init_q ? :EncoderInitQ : :EncoderInitP, + decoder_init_q ? :DecoderInitQ : :DecoderInitP, typeof(tanh)} + end +end + +@testset "ResNetLayer without and with a bias ($T)" for T in REAL_ELTYPES + Random.seed!(1) + l = ResNetLayer(3, tanh; use_bias = false) + ps = params(NeuralNetwork(Chain(l), CPU(), T))[1] + @test keys(ps) == (:weight,) + @test parameterlength(l) == 9 == sum(length, values(ps)) + x = T[1 2; -1 0; 0.5 3] + y = l(x, ps) + @test eltype(y) === T + @test y ≈ x + tanh.(ps.weight * x) + + l = ResNetLayer(3, tanh; use_bias = true) + ps = params(NeuralNetwork(Chain(l), CPU(), T))[1] + @test keys(ps) == (:weight, :bias) + @test parameterlength(l) == 12 == sum(length, values(ps)) + ps = (weight = ps.weight, bias = T[1, -1, 2]) + y = l(x, ps) + @test eltype(y) === T + @test y ≈ x + tanh.(ps.weight * x .+ ps.bias) +end + +@testset "the lower and the upper layer are one struct" begin + @test typeof(VolumePreservingLowerLayer(3)).name === + typeof(VolumePreservingUpperLayer(3)).name + @test VolumePreservingLowerLayer(3) isa VolumePreservingFeedForwardLayer + @test VolumePreservingUpperLayer(3) isa VolumePreservingFeedForwardLayer + @test length(methods_in(initialparameters, + joinpath("layers", "volume_preserving_feedforward.jl"))) == 1 +end + +@testset "the bias keyword is init_bias in every layer ($T)" for T in REAL_ELTYPES + rng = Random.default_rng() + for l in (VolumePreservingLowerLayer(3), VolumePreservingUpperLayer(3), ResNetLayer(3)) + ps = initialparameters( + rng, ZeroInitializer(), l, CPU(), T; init_bias = OneInitializer()) + @test ps.bias == ones(T, 3) + @test eltype(ps.bias) === T + end + weight = initialparameters( + rng, ZeroInitializer(), VolumePreservingLowerLayer(3), CPU(), + T).weight + @test weight isa StrictlyLowerTriangular{T} + weight = initialparameters( + rng, ZeroInitializer(), VolumePreservingUpperLayer(3), CPU(), + T).weight + @test weight isa StrictlyUpperTriangular{T} +end + +@testset "MultiHeadAttention has one initialparameters" begin + @test length(methods_in(initialparameters, joinpath("layers", "multi_head_attention.jl"))) == + 1 +end + +@testset "MultiHeadAttention draws its weights head by head ($T)" for T in REAL_ELTYPES + # The parameters of two heads are those of one head drawn twice in a row from the same stream: + # for each head the query, the key and the value, in this order. + rng = Random.Xoshiro(7) + ps = initialparameters(rng, GML.AbstractNeuralNetworks.GlorotUniform(), + MultiHeadAttention(4, 2), CPU(), T) + rng = Random.Xoshiro(7) + one_head = GML.AbstractNeuralNetworks.GlorotUniform() + expected = [begin + W = zeros(T, 4, 2) + one_head(rng, W) + W + end + for _ in 1:6] + @test [ps.PQ.head_1, ps.PK.head_1, ps.PV.head_1, + ps.PQ.head_2, ps.PK.head_2, ps.PV.head_2] == + expected + @test eltype(ps.PQ.head_1) === T +end + +@testset "MultiHeadAttention draws Stiefel weights head by head ($T)" for T in REAL_ELTYPES + # The same order with `Stiefel = true`, where each weight is a point of the Stiefel manifold + # that `StiefelLayer` draws. + rng = Random.Xoshiro(7) + ps = initialparameters(rng, GML.AbstractNeuralNetworks.GlorotUniform(), + MultiHeadAttention(4, 2; Stiefel = true), CPU(), T) + rng = Random.Xoshiro(7) + expected = [initialparameters(rng, GML.AbstractNeuralNetworks.GlorotUniform(), + GML.StiefelLayer{4, 2}(), CPU(), T).weight + for _ in 1:6] + weights = [ps.PQ.head_1, ps.PK.head_1, ps.PV.head_1, + ps.PQ.head_2, ps.PK.head_2, ps.PV.head_2] + @test all(W -> W isa StiefelManifold{T}, weights) + @test [W.A for W in weights] == [W.A for W in expected] + @test all(W -> W.A' * W.A ≈ I(2), weights) +end + +@testset "MultiHeadAttention finds each head by its name, $(ndims(x)) dimensions ($T)" for T in REAL_ELTYPES, + x in (T[sin(i + 2j) for i in 1:6, j in 1:3], T[sin(i + 2j + 3k) + for i in 1:6, j in 1:3, k in 1:2]) + + Random.seed!(3) + d = MultiHeadAttention(6, 3) + ps = params(NeuralNetwork(Chain(d), CPU(), T))[1] + reordered(nt) = (head_3 = nt.head_3, head_1 = nt.head_1, head_2 = nt.head_2) + y = d(x, ps) + @test eltype(y) === T + @test d(x, (PQ = reordered(ps.PQ), PK = ps.PK, PV = ps.PV)) == y + @test d(x, (PQ = reordered(ps.PQ), PK = reordered(ps.PK), PV = reordered(ps.PV))) == y +end + +@testset "every type argument in src/layers/ is spelled ::Type{T}" begin + hits = String[] + for file in readdir(joinpath(pkgdir(GML), "src", "layers"); join = true) + for (n, line) in enumerate(eachline(file)) + occursin(r"T::Type\b", line) && push!(hits, "$(basename(file)):$n") + end + end + @test hits == String[] +end diff --git a/test/layers/multi_head_attention_allocation.jl b/test/layers/multi_head_attention_allocation.jl new file mode 100644 index 000000000..b87e190e8 --- /dev/null +++ b/test/layers/multi_head_attention_allocation.jl @@ -0,0 +1,74 @@ +# The forward pass of `MultiHeadAttention` allocates what the same head products allocate when +# their keys are constants: it builds no head key per pass. `reference_output` below is the forward +# pass with the key of head `i` read from a tuple that is built once, outside the measurement. +# +# One rebuilt key costs 192 B a pass (Julia 1.13.2, one head), so k = 1000 passes put a key built +# per pass 192 000 B above the reference, 187 times the tolerance of 1024 B. The tolerance absorbs +# the 16-byte quantisation of the figures on some platforms. + +using GeometricMachineLearning +using GeometricMachineLearning: params, _mul, tensor_transpose_tensor_mul, tensor_tensor_mul +using Test +import Random + +include("../helpers/eltypes.jl") + +const PASSES = 1000 +const TOLERANCE = 1024 + +function reference_output(d::MultiHeadAttention, x::AbstractMatrix{T}, ps, head_keys) where {T} + dim = size(x, 1) + head_outputs = map(1:(d.n_heads)) do i + key = head_keys[i] + _mul(ps.PV[key]', x) * + d.activation(_mul(ps.PQ[key]', x)' * _mul(ps.PK[key]', x) / T(sqrt(dim))) + end + x + reduce(vcat, head_outputs) +end + +function reference_output(d::MultiHeadAttention, x::AbstractArray{T, 3}, ps, head_keys) where {T} + dim = size(x, 1) + head_outputs = map(1:(d.n_heads)) do i + key = head_keys[i] + 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)))) + end + x + vcat(head_outputs...)::eltype(head_outputs) +end + +# The bytes that `k` calls of `f(d, x, ps)` allocate, after one call that compiles it. +function allocations(f::F, d, x, ps, k::Integer) where {F} + f(d, x, ps) + @allocated for _ in 1:k + f(d, x, ps) + end +end + +forward(d, x, ps) = d(x, ps) + +struct WithKeys{K} + head_keys::K +end +(r::WithKeys)(d, x, ps) = reference_output(d, x, ps, r.head_keys) + +@testset "MultiHeadAttention builds no key per pass, n_heads = $n_heads, $(ndims(x₀))-array ($T)" for T in REAL_ELTYPES, + n_heads in (1, 2), + x₀ in (zeros(4, 3), zeros(4, 3, 2)) + + Random.seed!(n_heads) + d = MultiHeadAttention(4, n_heads) + ps = params(NeuralNetwork(Chain(d), CPU(), T))[1] + x = T.(x₀ .+ reshape(1:length(x₀), size(x₀)) ./ length(x₀)) + reference = WithKeys(ntuple(i -> Symbol("head_", i), n_heads)) + + @test reference(d, x, ps) == forward(d, x, ps) + @test eltype(forward(d, x, ps)) === T + + literal = allocations(reference, d, x, ps, PASSES) + measured = allocations(forward, d, x, ps, PASSES) + @test literal > 0 + @test abs(measured - literal) ≤ TOLERANCE +end diff --git a/test/layers/multi_head_attention_inference.jl b/test/layers/multi_head_attention_inference.jl index 9e2e4009f..f122977e7 100644 --- a/test/layers/multi_head_attention_inference.jl +++ b/test/layers/multi_head_attention_inference.jl @@ -1,10 +1,11 @@ # Both `compute_output_of_mha` methods concatenate the head outputs in one call. The matrix method # reaches Base's linear `vcat` through `reduce`; the tensor method cannot, because `reduce(vcat, …)` # folds pairwise on a 3-tensor, so it splats -- and a splatted vector hides the argument count from -# inference, which is why that call carries a result-type assertion. Nothing else in `test/` pins -# the inference of a forward pass, so an inference regression in either method is invisible to the -# suite: the layer still returns the right numbers, and the `Any` propagates out through the functor -# into everything downstream. See `CHANGELOG.md` for the change this guards. +# inference, which is why that call carries a result-type assertion. An inference regression in +# either method returns the right numbers, and the `Any` propagates out through the functor into +# everything downstream, so this file pins the inference of both methods, at 4 heads and at 32: a +# `map` over a tuple of more than 31 heads no longer infers. `test/architectures/forward_inference.jl` +# pins the forward pass of each architecture. using GeometricMachineLearning using Test @@ -17,8 +18,11 @@ Random.seed!(1234) mha_parameters(d, T) = GML.params(NeuralNetwork( Chain(d), KernelAbstractions.CPU(), T))[1] -@testset "compute_output_of_mha infers concretely, Stiefel = $Stiefel" for Stiefel in (false, true) - dim, n_heads, seq_length, n_data = 8, 4, 6, 3 +@testset "compute_output_of_mha infers concretely, $n_heads heads, Stiefel = $Stiefel" for n_heads in ( + 4, 32), + Stiefel in (false, true) + + dim, seq_length, n_data = 2n_heads, 6, 3 d = MultiHeadAttention(dim, n_heads; Stiefel = Stiefel) ps = mha_parameters(d, Float32) diff --git a/test/layers/positional_encoding.jl b/test/layers/positional_encoding.jl index e667cee6d..2258480f3 100644 --- a/test/layers/positional_encoding.jl +++ b/test/layers/positional_encoding.jl @@ -3,6 +3,8 @@ using GeometricMachineLearning using GeometricMachineLearning: initialparameters using AbstractNeuralNetworks: parameterlength using KernelAbstractions +using JLArrays: JLArray, JLBackend +import GPUArraysCore import Random import Zygote @@ -131,8 +133,8 @@ end end @testset "a network carrying the layer can be trained" begin - # The encoding is a constant, built with a mutating loop. Zygote refuses to differentiate that - # loop unless the builder is declared non-differentiable. Without the declaration every + # The encoding is a constant, built with a mutating broadcast. Zygote refuses to differentiate + # that broadcast unless the builder is declared non-differentiable. Without the declaration every # gradient through a network with `positional_encoding = true` raises # `Mutating arrays is not supported`. A layer that cannot be trained through is not a layer. x = rand(Float32, 4, 5) @@ -160,6 +162,31 @@ end end end +@testset "the layer runs on a JLArray with scalar indexing disallowed ($T)" for T in (Float32, Float64) + x = T[sin(i + 2j) for i in 1:4, j in 1:3] + x_device = JLArray(x) + l = PositionalEncoding(4) + GPUArraysCore.@allowscalar begin + task_local_storage(:ScalarIndexing, GPUArraysCore.ScalarDisallowed) do + # Scalar indexing is off in this block, so a host loop over the device array raises. + @test_throws "Scalar indexing is disallowed" x_device[1] + # A broadcast on a JLArray also reads a host matrix, so the output alone cannot show + # where the encoding lives: the encoding itself must be a JLArray. + P = GeometricMachineLearning._positional_encoding_like(x_device, 4) + @test P isa JLArray{T, 2} + @test Array(P) == positional_encoding(T, 4, 3) + y = l(x_device, NamedTuple()) + @test y isa JLArray{T, 2} + @test KernelAbstractions.get_backend(y) isa JLBackend + @test Array(y) == l(x, NamedTuple()) + + y = l(JLArray(reshape(x, 4, 3, 1)), NamedTuple()) + @test y isa JLArray{T, 3} + @test Array(y)[:, :, 1] == l(x, NamedTuple()) + end + end +end + @testset "the Transformer keyword puts the layer at the front, and only then" begin with_encoding = Transformer(4, 2, 2; positional_encoding = true) without = Transformer(4, 2, 2) diff --git a/test/runtests.jl b/test/runtests.jl index aaf817dd1..ca1821161 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -35,6 +35,9 @@ if "core" in GROUPS @safetestset "Linear Symplectic Attention" include("layers/linear_symplectic_attention.jl") @safetestset "parameterlength(::MultiHeadAttention{M,M,true}) is exact" include("layers/multi_head_attention_parameterlength.jl") @safetestset "compute_output_of_mha infers concretely" include("layers/multi_head_attention_inference.jl") + @safetestset "MultiHeadAttention builds no head key per forward pass" include("layers/multi_head_attention_allocation.jl") + @safetestset "Layer forward passes against the recorded outputs" include("layers/layer_pins.jl") + @safetestset "Layer type parameters, shared methods and type arguments" include("layers/layer_structure.jl") @safetestset "Documented keyword defaults are the constructed ones" include("integration/documented_defaults.jl") @safetestset "Test setup of transformer with Stiefel weights" include("layers/transformer_setup.jl") @safetestset "Sinusoidal positional encoding" include("layers/positional_encoding.jl") @@ -50,6 +53,9 @@ if "core" in GROUPS @safetestset "PSD tests" include("architectures/psd_architecture_tests.jl") @safetestset "SymplecticAutoencoder tests" include("architectures/symplectic_autoencoder_tests.jl") @safetestset "Check if autoencoder error is lower than PSD error" include("architectures/sae_error_lower_than_psd_error.jl") + @safetestset "Architecture forward passes against the recorded outputs" include("architectures/architecture_pins.jl") + @safetestset "Architecture keywords and chain builders" include("architectures/architecture_structure.jl") + @safetestset "The forward pass of every exported architecture infers" include("architectures/forward_inference.jl") @safetestset "Check reduced model" include("reduced_system/reduced_system.jl") @safetestset "_norm keeps the element type of its argument" include("utils.jl") @safetestset "Symplectic Euler and variational midpoint losses" include("loss/training_method_losses.jl")