diff --git a/Project.toml b/Project.toml index f8031e5b..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" @@ -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"] 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)) 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(;