From ce48aa5df1ffde0a4dac8e0aa33967337d0eab93 Mon Sep 17 00:00:00 2001 From: ChrisRackauckas-Claude Date: Wed, 8 Jul 2026 05:54:03 -0400 Subject: [PATCH] Document exported public API Co-Authored-By: Chris Rackauckas --- docs/src/api.md | 2 + src/layers.jl | 165 +++++++++++++++++++++++++++++++----------------- 2 files changed, 110 insertions(+), 57 deletions(-) diff --git a/docs/src/api.md b/docs/src/api.md index 0a2384db..9a1ee427 100644 --- a/docs/src/api.md +++ b/docs/src/api.md @@ -49,6 +49,8 @@ recommend: ## Standard Models +`DEQs` is exported as an alias for the `DeepEquilibriumNetworks` module. + ```@docs DeepEquilibriumNetwork SkipDeepEquilibriumNetwork diff --git a/src/layers.jl b/src/layers.jl index ce8aa42d..b6f2fb5b 100644 --- a/src/layers.jl +++ b/src/layers.jl @@ -1,17 +1,16 @@ """ - DeepEquilibriumSolution(z_star, uâ‚€, residual, jacobian_loss, nfe, solution) + DeepEquilibriumSolution(z_star, u0, residual, jacobian_loss, nfe, original) -Stores the solution of a DeepEquilibriumNetwork and its variants. +Stores the solution data produced by a [`DeepEquilibriumNetwork`](@ref) and its variants. ## Fields - - `z_star`: Steady-State or the value reached due to maxiters - - `u0`: Initial Condition - - `residual`: Difference of the ``z^*`` and ``f(z^*, x)`` - - `jacobian_loss`: Jacobian Stabilization Loss (see individual networks to see how it - can be computed) - - `nfe`: Number of Function Evaluations - - `original`: Original Internal Solution + - `z_star`: Steady state, or the final iterate reached when the solver stops. + - `u0`: Initial condition used by the equilibrium solve. + - `residual`: Difference between ``z^*`` and ``f(z^*, x)``. + - `jacobian_loss`: Jacobian stabilization loss. + - `nfe`: Number of function evaluations. + - `original`: Original solver solution object. """ struct DeepEquilibriumSolution # This is intentionally left untyped to allow updating `st` z_star @@ -72,6 +71,54 @@ function Base.show(io::IO, sol::DeepEquilibriumSolution) end # Core Model +""" + DeepEquilibriumNetwork(model, solver; init = missing, jacobian_regularization = nothing, + problem_type::Type = SteadyStateProblem{false}, kwargs...) + +Deep Equilibrium Network as proposed in [baideep2019](@cite) and [pal2022mixing](@cite). + +## Arguments + + - `model`: Lux layer defining the equilibrium map. + - `solver`: Solver for the equilibrium problem. ODE solvers and nonlinear solvers are + both supported. + +## Keywords + + - `init`: Initial condition layer for the equilibrium problem. If `nothing`, the initial + condition is set to `zero(x)`. If `missing`, it is set to `WrappedFunction(zero)`. + Otherwise, pass a Lux layer called as `init(x, ps, st)`. + - `jacobian_regularization`: Jacobian stabilization backend. Supported values are + `nothing`, `AutoForwardDiff`, `AutoFiniteDiff`, and `AutoZygote`. + - `problem_type`: Equilibrium problem type. Use `ODEProblem` to construct an ODE-based + network; defaults to `SteadyStateProblem`. + - `kwargs`: Additional keyword arguments passed to `SciMLBase.solve`. + +## Returns + +Returns a Lux layer. Calling the layer returns the model output and a state whose `solution` +field contains a [`DeepEquilibriumSolution`](@ref). + +## Example + +```jldoctest +julia> using DeepEquilibriumNetworks, Lux, SteadyStateDiffEq, Random + +julia> model = DeepEquilibriumNetwork( + Parallel(+, Dense(2, 2; use_bias=false), Dense(2, 2; use_bias=false)), + SSRootfind(); verbose=false); + +julia> rng = Xoshiro(0); + +julia> ps, st = Lux.setup(rng, model); + +julia> size(first(model(ones(Float32, 2, 1), ps, st))) +(2, 1) +``` + +See also: [`SkipDeepEquilibriumNetwork`](@ref), [`MultiScaleDeepEquilibriumNetwork`](@ref), +[`MultiScaleSkipDeepEquilibriumNetwork`](@ref). +""" @concrete struct DeepEquilibriumNetwork <: AbstractLuxContainerLayer{(:model, :init)} init model @@ -161,51 +208,6 @@ function (deq::DEQ)(x, ps, st::NamedTuple, ::Val{false}) end ## Constructors -""" - DeepEquilibriumNetwork(model, solver; init = missing, jacobian_regularization=nothing, - problem_type::Type=SteadyStateProblem{false}, kwargs...) - -Deep Equilibrium Network as proposed in [baideep2019](@cite) and [pal2022mixing](@cite). - -## Arguments - - - `model`: Neural Network. - - `solver`: Solver for the rootfinding problem. ODE Solvers and Nonlinear Solvers are both - supported. - -## Keyword Arguments - - - `init`: Initial Condition for the rootfinding problem. If `nothing`, the initial - condition is set to `zero(x)`. If `missing`, the initial condition is set to - `WrappedFunction(zero)`. In other cases the initial condition is set to - `init(x, ps, st)`. - - `jacobian_regularization`: Must be one of `nothing`, `AutoForwardDiff`, `AutoFiniteDiff` - or `AutoZygote`. - - `problem_type`: Provides a way to simulate a Vanilla Neural ODE by setting the - `problem_type` to `ODEProblem`. By default, the problem type is set to - `SteadyStateProblem`. - - `kwargs`: Additional Parameters that are directly passed to `SciMLBase.solve`. - -## Example - -```jldoctest -julia> using DeepEquilibriumNetworks, Lux, SteadyStateDiffEq, Random - -julia> model = DeepEquilibriumNetwork( - Parallel(+, Dense(2, 2; use_bias=false), Dense(2, 2; use_bias=false)), - SSRootfind(); verbose=false); - -julia> rng = Xoshiro(0); - -julia> ps, st = Lux.setup(rng, model); - -julia> size(first(model(ones(Float32, 2, 1), ps, st))) -(2, 1) -``` - -See also: [`SkipDeepEquilibriumNetwork`](@ref), [`MultiScaleDeepEquilibriumNetwork`](@ref), -[`MultiScaleSkipDeepEquilibriumNetwork`](@ref). -""" function DeepEquilibriumNetwork( model, solver; init = missing, jacobian_regularization = nothing, problem_type::Type = SteadyStateProblem{false}, kwargs... @@ -226,8 +228,36 @@ end """ SkipDeepEquilibriumNetwork(model, [init=nothing,] solver; kwargs...) -Skip Deep Equilibrium Network as proposed in [pal2022mixing](@cite). Alias which creates -a [`DeepEquilibriumNetwork`](@ref) with `init` kwarg set to passed value. +Skip Deep Equilibrium Network as proposed in [pal2022mixing](@cite). + +This is a convenience constructor for [`DeepEquilibriumNetwork`](@ref) that forwards `init` +through the `init` keyword argument. If `init` is omitted, the initial condition is +`nothing`. + +## Arguments + + - `model`: Lux layer defining the equilibrium map. + - `init`: Optional Lux layer used to construct the initial condition. + - `solver`: Solver for the equilibrium problem. + +## Returns + +Returns a [`DeepEquilibriumNetwork`](@ref). + +## Example + +```jldoctest +julia> using DeepEquilibriumNetworks, Lux, SteadyStateDiffEq, Random + +julia> model = SkipDeepEquilibriumNetwork( + Parallel(+, Dense(2, 2; use_bias=false), Dense(2, 2; use_bias=false)), + SSRootfind(); verbose=false); + +julia> ps, st = Lux.setup(Xoshiro(0), model); + +julia> size(first(model(ones(Float32, 2, 1), ps, st))) +(2, 1) +``` """ function SkipDeepEquilibriumNetwork(model, init, solver; kwargs...) return DeepEquilibriumNetwork(model, solver; init, kwargs...) @@ -312,6 +342,19 @@ Skip Multi Scale Deep Equilibrium Network as proposed in [pal2022mixing](@cite). creates a [`MultiScaleDeepEquilibriumNetwork`](@ref) with `init` kwarg set to passed value. If `init` is not passed, it creates a MultiScale Regularized Deep Equilibrium Network. + +## Arguments + + - `main_layers`: Tuple of Lux layers, one per scale. + - `mapping_layers`: Matrix of Lux layers mapping between scales. + - `post_fuse_layer`: Optional tuple of Lux layers applied after scale fusion. + - `init`: Optional tuple of Lux layers used to construct the initial conditions. + - `solver`: Solver for the equilibrium problem. + - `scales`: Output shape for each scale. + +## Returns + +Returns a [`MultiScaleDeepEquilibriumNetwork`](@ref). """ function MultiScaleSkipDeepEquilibriumNetwork( main_layers::Tuple, mapping_layers::Matrix, @@ -335,8 +378,16 @@ end """ MultiScaleNeuralODE(args...; kwargs...) -Same arguments as [`MultiScaleDeepEquilibriumNetwork`](@ref) but sets `problem_type` to +Construct a multi-scale neural ODE with the same arguments as +[`MultiScaleDeepEquilibriumNetwork`](@ref). + +This forwards all positional and keyword arguments to +[`MultiScaleDeepEquilibriumNetwork`](@ref), while setting `problem_type` to `ODEProblem{false}`. + +## Returns + +Returns a [`MultiScaleDeepEquilibriumNetwork`](@ref) configured with ODE dynamics. """ function MultiScaleNeuralODE(args...; kwargs...) return MultiScaleDeepEquilibriumNetwork(args...; kwargs..., problem_type = ODEProblem{false})