Quick start example fails to run when batch size increases moderately

[AnasAbdelR]

To reproduce:

using DeepEquilibriumNetworks, Lux, Random, Zygote
# using LuxCUDA, LuxAMDGPU ## Install and Load for GPU Support

seed = 0
rng = Random.default_rng()
Random.seed!(rng, seed)
model = Chain(Dense(2 => 2),
    DeepEquilibriumNetwork(Parallel(+,
            Dense(2 => 2; use_bias=false),
            Dense(2 => 2; use_bias=false)),
        ContinuousDEQSolver(; abstol=0.1f0, reltol=0.1f0, abstol_termination=0.1f0,
            reltol_termination=0.1f0);
        save_everystep=true))

gdev = gpu_device()
cdev = cpu_device()

ps, st = Lux.setup(rng, model) |> gdev
x = rand(rng, Float32, 2, 100) |> gdev
y = rand(rng, Float32, 2, 100) |> gdev
gs = only(Zygote.gradient(p -> sum(abs2, first(first(model(x, p, st))) .- y), ps))

This gives the following error + warning in 1.9 (used a try,catch log because the original error flooded the repl and can't be accessed 😅 ):

┌ Warning: Automatic AD choice of autojacvec failed in ODE adjoint, failing back to ODE adjoint + numerical vjp
└ @ SciMLSensitivity ~/.julia/packages/SciMLSensitivity/U8Axh/src/sensitivity_interface.jl:381
┌ Warning: AD choice of autojacvec failed in nonlinear solve adjoint
└ @ SciMLSensitivity ~/.julia/packages/SciMLSensitivity/U8Axh/src/steadystate_adjoint.jl:112
1278

and

Error encountered: MethodError: no method matching jacobian(::SciMLSensitivity.ParamGradientWrapper{ODEFunction{false, SciMLBase.FullSpecialize, DeepEquilibriumNetworks.var"#dudt#50"{Lux.Experimental.StatefulLuxLayer{Parallel{NamedTuple{(:layer_1, :layer_2), Tuple{Dense{false, typeof(identity), typeof(glorot_uniform), typeof(zeros32)}, Dense{false, typeof(identity), typeof(glorot_uniform), typeof(zeros32)}}}, Nothing, typeof(+)}, Nothing, NamedTuple{(:layer_1, :layer_2), Tuple{NamedTuple{(), Tuple{}}, NamedTuple{(), Tuple{}}}}}}, LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Nothing, Matrix{Float32}}, ::NamedTuple{(:ps, :x), Tuple{NamedTuple{(:layer_1, :layer_2), Tuple{NamedTuple{(:weight,), Tuple{Matrix{Float32}}}, NamedTuple{(:weight,), Tuple{Matrix{Float32}}}}}, Matrix{Float32}}}, ::SteadyStateAdjoint{0, true, Val{:central}, Bool, Nothing, NamedTuple{(), Tuple{}}})

Closest candidates are:
  jacobian(::Any, !Matched::AbstractArray{<:Number}, ::SciMLBase.AbstractOverloadingSensitivityAlgorithm)
   @ SciMLSensitivity ~/.julia/packages/SciMLSensitivity/U8Axh/src/derivative_wrappers.jl:128

[avik-pal]

Use componentarrays. But this needs to be addressed in a more general setup, but I won't have time to do it rn