Update deps

Project.toml

@@ -6,6 +6,7 @@ version = "0.1.0"
 [deps]
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
+CoupledNODE = "88291d29-22ea-41b1-bc0b-03785bffce48"
 Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
 LuxCore = "bb33d45b-7691-41d6-9220-0943567d0623"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"

test/test-model.jl

@@ -11,10 +11,10 @@ using Zygote: Zygote
     # Define parameters for the model
     T = Float32
     D = 2
-    N = 64       # Define the spatial dimension for the attention layer input
-    emb_size = 32
+    N = 32       # Define the spatial dimension for the attention layer input
+    emb_size = 16
     patch_size = 4
-    n_heads = 4
+    n_heads = 2
     d = emb_size ÷ n_heads  # Dimension per attention head
     rng = Xoshiro(123)
 
@@ -32,11 +32,10 @@ using Zygote: Zygote
 
     # Verify CNN layers' setup
     @test CnnLayers != nothing
-    @test typeof(CnnLayers) <: Lux.AbstractLayer
 
     # Define the Transformer model layers
-    attention_layer = AttentionLayer(N, d, emb_size, patch_size, n_heads; T = T)
-    @test typeof(attention_layer) <: Lux.AbstractLayer
+    attention_layer = attention(N, D, emb_size, patch_size, n_heads; T = T)
+    @test attention_layer != nothing
 
     # Combine Attention Layer with CNN in a Lux chain
     layers = (
@@ -60,65 +59,11 @@ using Zygote: Zygote
     output = closure(input_tensor, θ, st)
 
     # Test output dimensions after passing through the model
-    expected_channels = 2  # From the CNN output
+    expected_channels = D  # From the CNN output
     @test size(output[1]) == (N, N, expected_channels, 1)  # Check final output size
 
     # Test Differentiability by calculating gradients
     grad = Zygote.gradient(θ -> sum(abs2, closure(input_tensor, θ, st)[1]), θ)
     @test !isnothing(grad)  # Ensure gradients were successfully computed
 
 end
-
-# Define parameters for the model
-T = Float32
-D = 2
-N = 64       # Define the spatial dimension for the attention layer input
-emb_size = 32
-patch_size = 4
-n_heads = 4
-d = emb_size ÷ n_heads  # Dimension per attention head
-rng = Xoshiro(123)
-
-# Test CNN Layer Setup
-CnnLayers, _, _ = cnn(;
-    T = T,
-    D = D,
-    data_ch = 2 * D,                     # Input channels for CNN after concatenation
-    radii = [3, 3],
-    channels = [2, 2],
-    activations = [tanh, identity],
-    use_bias = [false, false],
-    rng,
-)
-
-# Verify CNN layers' setup
-@test CnnLayers != nothing
-
-# Define the Transformer model layers
-attention_layer = attention(N, d, emb_size, patch_size, n_heads; T = T)
-@test typeof(attention_layer) <: Lux.AbstractLayer
-
-# Combine Attention Layer with CNN in a Lux chain
-layers = (
-    Lux.SkipConnection(attention_layer, (x, y) -> cat(x, y; dims = 3); name = "Attention"),
-    CnnLayers,
-)
-closure = Lux.Chain(layers...)
-θ, st = Lux.setup(rng, closure)
-θ = ComponentArray(θ)
-
-# Test model structure
-@test typeof(closure) <: Lux.Chain
-@test length(closure.layers) == 2  # Confirm both layers (Attention + CNN) are in chain
-
-# Define input tensor and pass through model
-input_tensor = rand(T, N, N, D, 1)  # Example input with shape (N, N, D, batch_size)
-output = closure(input_tensor, θ, st)
-
-# Test output dimensions after passing through the model
-expected_channels = 2  # From the CNN output
-@test size(output[1]) == (N, N, expected_channels, 1)  # Check final output size
-
-# Test Differentiability by calculating gradients
-grad = Zygote.gradient(θ -> sum(abs2, closure(input_tensor, θ, st)[1]), θ)
-@test !isnothing(grad)  # Ensure gradients were successfully computed