Fix the projection

Author: avik-pal

src/utils.jl

@@ -26,12 +26,10 @@ end
 
 function CRC.rrule(::typeof(__flatten_vcat), x)
     y = __flatten_vcat(x)
-    projects = CRC.ProjectTo.(x)
+    project_x = CRC.ProjectTo(x)
     function ∇__flatten_vcat(∂y)
         ∂y isa CRC.NoTangent && return (CRC.NoTangent(), CRC.NoTangent())
-        ∂x = __split_and_reshape(∂y, x)
-        ∂x = map((∂xᵢ, project) -> project(∂xᵢ), ∂x, projects)
-        return CRC.NoTangent(), ∂x
+        return CRC.NoTangent(), project_x(__split_and_reshape(∂y, x))
     end
     return y, ∇__flatten_vcat
 end

test/Project.toml

@@ -6,6 +6,7 @@ DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
 Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
+LuxCUDA = "d0bbae9a-e099-4d5b-a835-1c6931763bda"
 LuxTestUtils = "ac9de150-d08f-4546-94fb-7472b5760531"
 NLsolve = "2774e3e8-f4cf-5e23-947b-6d7e65073b56"
 NonlinearSolve = "8913a72c-1f9b-4ce2-8d82-65094dcecaec"
@@ -21,4 +22,4 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]
-Aqua = "0.8"
+Aqua = "0.8"

test/layers.jl

@@ -111,10 +111,8 @@ end
     jacobian_regularizations = (nothing, AutoFiniteDiff(), AutoZygote())
 
     for mtype in model_type, jacobian_regularization in jacobian_regularizations
-        # @testset "Solver: $(__nameof(solver))"
-        for solver in solvers
-            # @testset "x_size: $(x_size)"
-            for (main_layer, mapping_layer, init_layer, x_size, scale) in zip(main_layers,
+        @testset "Solver: $(__nameof(solver))" for solver in solvers
+            @testset "x_size: $(x_size)" for (main_layer, mapping_layer, init_layer, x_size, scale) in zip(main_layers,
                 mapping_layers, init_layers, x_sizes, scales)
                 @info solver, mtype, jacobian_regularization, main_layer, mapping_layer,
                 init_layer, x_size, scale