Fix Runic format check: add return in test/runtests.jl core_body()

Project.toml

@@ -1,7 +1,7 @@
 name = "DeepEquilibriumNetworks"
 uuid = "6748aba7-0e9b-415e-a410-ae3cc0ecb334"
 authors = ["Avik Pal <avikpal@mit.edu>"]
-version = "2.7.0"
+version = "2.7.1"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -37,6 +37,7 @@ GPUArraysCore = "0.1, 0.2"
 SafeTestsets = "0.1"
 InteractiveUtils = "<0.0.1, 1"
 Lux = "1"
+LuxCUDA = "0.3"
 LuxCore = "1"
 LuxTestUtils = "1, 2"
 MLDataDevices = "1"
@@ -67,6 +68,7 @@ Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
 GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
+LuxCUDA = "d0bbae9a-e099-4d5b-a835-1c6931763bda"
 LuxTestUtils = "ac9de150-d08f-4546-94fb-7472b5760531"
 MLDataDevices = "7e8f7934-dd98-4c1a-8fe8-92b47a384d40"
 NLsolve = "2774e3e8-f4cf-5e23-947b-6d7e65073b56"
@@ -80,4 +82,4 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["Aqua", "Documenter", "ExplicitImports", "ForwardDiff", "Functors", "GPUArraysCore", "InteractiveUtils", "LuxTestUtils", "MLDataDevices", "NLsolve", "NonlinearSolve", "OrdinaryDiffEq", "OrdinaryDiffEqAdamsBashforthMoulton", "SafeTestsets", "SciMLSensitivity", "SciMLTesting", "StableRNGs", "Test", "Zygote"]
+test = ["Aqua", "Documenter", "ExplicitImports", "ForwardDiff", "Functors", "GPUArraysCore", "InteractiveUtils", "LuxCUDA", "LuxTestUtils", "MLDataDevices", "NLsolve", "NonlinearSolve", "OrdinaryDiffEq", "OrdinaryDiffEqAdamsBashforthMoulton", "SafeTestsets", "SciMLSensitivity", "SciMLTesting", "StableRNGs", "Test", "Zygote"]

docs/Project.toml

@@ -12,6 +12,7 @@ NonlinearSolve = "8913a72c-1f9b-4ce2-8d82-65094dcecaec"
 OneHotArrays = "0b1bfda6-eb8a-41d2-88d8-f5af5cad476f"
 Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
 OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
+OrdinaryDiffEqAdamsBashforthMoulton = "89bda076-bce5-4f1c-845f-551c83cdda9a"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SciMLSensitivity = "1ed8b502-d754-442c-8d5d-10ac956f44a1"
@@ -31,6 +32,7 @@ NonlinearSolve = "3, 4"
 OneHotArrays = "0.2"
 Optimisers = "0.3, 0.4"
 OrdinaryDiffEq = "6, 7"
+OrdinaryDiffEqAdamsBashforthMoulton = "1, 2"
 Random = "1"
 SciMLSensitivity = "7"
 Setfield = "1"

docs/src/tutorials/basic_mnist_deq.md

@@ -5,6 +5,7 @@ We will train a simple Deep Equilibrium Model on MNIST. First we load a few pack
 ```@example basic_mnist_deq
 using DeepEquilibriumNetworks, SciMLSensitivity, Lux, NonlinearSolve, OrdinaryDiffEq,
       Random, Optimisers, Zygote, LinearSolve, Dates, Printf, Setfield, OneHotArrays
+using OrdinaryDiffEqAdamsBashforthMoulton: VCAB3
 using MLDatasets: MNIST
 using MLUtils: DataLoader, splitobs
 using LuxCUDA # For NVIDIA GPU support

src/layers.jl

@@ -128,7 +128,11 @@ function (deq::DEQ)(x, ps, st::NamedTuple, ::Val{false})
         # The type-assert is needed because of an upstream Lux issue with type stability of
         # conv with Dual numbers
         y = model((u, p.x), p.ps)::typeof(u)
-        return y .- u
+        # Subtract in flattened space: the matrix-free SteadyStateAdjoint VJP seeds the
+        # residual pullback with a flat cotangent (it builds the operator from `vec(y)`),
+        # which Zygote's `ProjectTo` rejects when broadcasting against a multi-dim `u`
+        # (e.g. a conv state). `vec`/`reshape` round-trip is a no-op on the value.
+        return reshape(vec(y) .- vec(u), size(u))
     end
 
     prob = construct_prob(deq.kind, ODEFunction{false}(dudt), z, (; ps = ps.model, x))

test/runtests.jl

@@ -7,7 +7,7 @@ using SafeTestsets
 # cannot be expressed as a separate folder of files. Hence explicit-args run_tests.
 function core_body()
     @safetestset "Utils Tests" include("utils_tests.jl")
-    @safetestset "Layers Tests" include("layers_tests.jl")
+    return @safetestset "Layers Tests" include("layers_tests.jl")
 end
 
 run_tests(;