From f3af78a09fbb5a9d6c7cec7af2147fbd0cdfbc03 Mon Sep 17 00:00:00 2001 From: ChrisRackauckas-Claude Date: Mon, 15 Jun 2026 10:09:02 -0400 Subject: [PATCH 1/3] Apply Runic formatting to test/runtests.jl The Runic Format Check on main fails because core_body()'s final expression (the Layers Tests @safetestset) is missing the explicit return that Runic requires. Add it to make the format check pass. Co-Authored-By: Chris Rackauckas Co-Authored-By: Claude Opus 4.8 (1M context) --- test/runtests.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/runtests.jl b/test/runtests.jl index e35b82b2..612f832e 100644 --- a/test/runtests.jl +++ b/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(; From 2c5e20ad8b581f3c21e7626c135199a298174daf Mon Sep 17 00:00:00 2001 From: ChrisRackauckas-Claude Date: Tue, 16 Jun 2026 05:31:49 -0400 Subject: [PATCH 2/3] Restore LuxCUDA test dependency dropped during v1.2 CI conversion The grouped-tests v1.2 conversion (#217) dropped LuxCUDA from the test [extras]/[targets]. The self-hosted GPU lane sets BACKEND_GROUP=CUDA, which makes test/shared_testsetup.jl do `using LuxCUDA`; without the dependency the GPU job errors with `ArgumentError: Package LuxCUDA not found in current path`. Re-add LuxCUDA (compat 0.3) to [compat], [extras], and the test target, using the correct registry UUID d0bbae9a-e099-4d5b-a835-1c6931763bda (the pre-conversion Project.toml had carried a typo'd UUID ...1c6931f17571). Co-Authored-By: Chris Rackauckas Co-Authored-By: Claude Opus 4.8 (1M context) --- Project.toml | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index f8031e5b..f3d0bf8a 100644 --- a/Project.toml +++ b/Project.toml @@ -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"] From e251de6f9863de6441363982e2efd0cba1140164 Mon Sep 17 00:00:00 2001 From: ChrisRackauckas-Claude Date: Tue, 16 Jun 2026 10:21:43 -0400 Subject: [PATCH 3/3] Fix MNIST docs build: steady-state adjoint VJP shape + missing VCAB3 import MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The Documentation build failed on two genuine bugs in the basic MNIST DEQ tutorial (both reproduce on `main`, independent of the test/Project.toml changes in this PR): 1. `DimensionMismatch` in the steady-state adjoint VJP. For a conv DEQ the rootfind state is a 4D array (e.g. `(13, 13, 64, 2)`). When the state has more than 50 elements, SciMLSensitivity's `SteadyStateAdjoint` takes the matrix-free `VecJacOperator` path and builds the operator from `vec(y)`, so it seeds the residual's pullback with a *flat* cotangent. Zygote's `ProjectTo` then rejects that flat cotangent when it is broadcast against the multi-dimensional `u` in the residual `y .- u`, raising `variable with size(x) == (13,13,64,2) cannot have a gradient with size(dx) == (21632,)`. The Core test suite never hit this because its conv states are <= 50 elements (dense-Jacobian path). Fix: compute the residual in flattened space, `reshape(vec(y) .- vec(u), size(u))` — a value-preserving no-op that keeps the broadcast operands consistent with the flat cotangent. 2. `UndefVarError: VCAB3`. The tutorial calls `VCAB3()` but only imports `OrdinaryDiffEq`, which (v7) no longer re-exports `VCAB3`. Added `OrdinaryDiffEqAdamsBashforthMoulton` to the docs project and an explicit `using OrdinaryDiffEqAdamsBashforthMoulton: VCAB3` in the tutorial. Verified locally on Julia 1.12: - Full `docs/make.jl` builds end-to-end (both MNIST `@example` blocks run; only a transient ColPrac linkcheck 429 from the sandbox differs from CI). - `GROUP=Core Pkg.test()` passes: Utils 11/11, Layers 1538/1538. Co-Authored-By: Chris Rackauckas Co-Authored-By: Claude Opus 4.8 (1M context) --- Project.toml | 2 +- docs/Project.toml | 2 ++ docs/src/tutorials/basic_mnist_deq.md | 1 + src/layers.jl | 6 +++++- 4 files changed, 9 insertions(+), 2 deletions(-) diff --git a/Project.toml b/Project.toml index f3d0bf8a..c89697d3 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "DeepEquilibriumNetworks" uuid = "6748aba7-0e9b-415e-a410-ae3cc0ecb334" authors = ["Avik Pal "] -version = "2.7.0" +version = "2.7.1" [deps] ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" diff --git a/docs/Project.toml b/docs/Project.toml index 718926c3..37c44735 100644 --- a/docs/Project.toml +++ b/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" diff --git a/docs/src/tutorials/basic_mnist_deq.md b/docs/src/tutorials/basic_mnist_deq.md index df2d4984..70f5108d 100644 --- a/docs/src/tutorials/basic_mnist_deq.md +++ b/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 diff --git a/src/layers.jl b/src/layers.jl index 263167b3..ce8aa42d 100644 --- a/src/layers.jl +++ b/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))