Better Steady State Adjoint

[avik-pal]

Sister PR to SciML/SciMLSensitivity.jl#877

- [ ] Compare the improved SSAdjoint Methods

Fast DEQ Solves using JNFK

using DeepEquilibriumNetworks, LuxAMDGPU, Lux, Random, OrdinaryDiffEq, Zygote,
    SciMLSensitivity, LinearAlgebra, NonlinearSolve, LegibleLambdas, Statistics, Optimisers,
    LinearSolve

import Lux.Experimental: @compact

function NonlinearSolve.SparseDiffTools.get_tag(::AbstractArray{
    NonlinearSolve.ForwardDiff.Dual{T, V, N}}) where {T, V, N}
    return T
end

dudt = @compact(; c1=CrossCor((3, 3), 1 => 4, swish; pad=SamePad()),
    c2=CrossCor((3, 3), 1 => 4, swish; pad=SamePad()),
    c3=Chain(;
        l1=CrossCor((3, 3), 8 => 32, swish; pad=SamePad()),
        l2=CrossCor((3, 3), 32 => 128, swish; pad=SamePad()),
        l3=CrossCor((3, 3), 128 => 32, swish; pad=SamePad()),
        l4=CrossCor((3, 3), 32 => 8, swish; pad=SamePad()),
        l5=CrossCor((3, 3), 8 => 1; pad=SamePad()))) do (u, x)
    return c3(cat(c1(u), c2(x); dims=Val(3)))
end

dev = gpu_device()

model = DeepEquilibriumNetwork(dudt,
    DiscreteDEQSolver(NewtonRaphson(; linsolve=SimpleGMRES()));
    abstol=1.0f-6, reltol=1.0f-6, maxiters=20)

x = randn(Float32, 32, 32, 1, 16) |> dev;

ps, st = Lux.setup(Xoshiro(0), model) |> dev;

model(x, ps, st)

[codecov]

Codecov Report

Merging #93 (34932dc) into main (30cdb93) will decrease coverage by 2.28%.
Report is 4 commits behind head on main.
The diff coverage is 95.00%.

@@            Coverage Diff             @@
##             main      #93      +/-   ##
==========================================
- Coverage   97.92%   95.65%   -2.28%     
==========================================
  Files           9        9              
  Lines         241      207      -34     
==========================================
- Hits          236      198      -38     
- Misses          5        9       +4     

Files	Coverage Δ
src/chainrules.jl	100.00% <100.00%> (ø)
src/layers/core.jl	100.00% <100.00%> (ø)
src/layers/deq.jl	100.00% <100.00%> (ø)
src/layers/evaluate.jl	97.22% <100.00%> (-2.78%)	⬇️
src/layers/jacobian_stabilization.jl	100.00% <100.00%> (+4.54%)	⬆️
src/layers/mdeq.jl	95.23% <100.00%> (-0.37%)	⬇️
src/solve.jl	100.00% <100.00%> (ø)
src/DeepEquilibriumNetworks.jl	33.33% <0.00%> (-66.67%)	⬇️
src/utils.jl	83.33% <50.00%> (-16.67%)	⬇️

📣 We’re building smart automated test selection to slash your CI/CD build times. Learn more

[avik-pal]

With SciML/SciMLSensitivity.jl#918 this should be done

[ChrisRackauckas]

Why is a special adjoint still needed?

[avik-pal]

There is no special adjoint. I just migrated to using the non-component arrays version

[avik-pal]

The name is from the SciMLSensitivity change 😅

[avik-pal]

Needs to wait till LuxDL/Lux.jl#425 is merged