Edifice.Utils.ODESolver (Edifice v0.2.0)

ODE Solver for Nx tensors - used by Liquid Neural Networks.

Provides numerical integration methods for solving ordinary differential equations of the form: dx/dt = f(t, x)

Available Solvers

Solver	Order	Adaptive	Best For
`:euler`	1	No	Fast, simple problems
`:midpoint`	2	No	Better accuracy than Euler
`:rk4`	4	No	Good accuracy, fixed step
`:dopri5`	4/5	Yes	Best accuracy, adaptive step

Usage

# Solve dx/dt = -x (exponential decay)
f = fn _t, x -> Nx.negate(x) end
x0 = Nx.tensor([1.0])
result = ODESolver.solve(f, 0.0, 1.0, x0, solver: :rk4, dt: 0.1)

Neural ODE Integration

For Liquid Neural Networks, use solve_ltc/5 which is optimized for the LTC (Liquid Time-Constant) ODE:

dx/dt = (-x + f(x, input)) / tau

Differentiability

All solvers use pure Nx operations and are compatible with Nx.Defn.grad/2 for automatic differentiation during training.

Reference

Dormand & Prince (1980): "A family of embedded Runge-Kutta formulae"
Chen et al. (2018): "Neural Ordinary Differential Equations"

Summary

Functions

solve(f, t0, t1, x0, opts \\ [])

Solve an ODE from t0 to t1 with initial condition x0.

solve_ltc(x, activation, tau, opts \\ [])

Solve the LTC (Liquid Time-Constant) ODE for one timestep.

Functions

solve(f, t0, t1, x0, opts \\ [])

@spec solve(
  (float(), Nx.Tensor.t() -> Nx.Tensor.t()),
  float(),
  float(),
  Nx.Tensor.t(),
  keyword()
) :: Nx.Tensor.t()

Solve an ODE from t0 to t1 with initial condition x0.

Parameters

f - The ODE function f(t, x) returning dx/dt
t0 - Initial time
t1 - Final time
x0 - Initial state (Nx tensor)
opts - Solver options

Options

:solver - Solver type: :euler, :midpoint, :rk4, :dopri5 (default: :rk4)
:dt - Time step for fixed-step methods (default: 0.01)
:max_steps - Maximum steps for adaptive methods (default: 1000)
:atol - Absolute tolerance for adaptive methods (default: 1.0e-6)
:rtol - Relative tolerance for adaptive methods (default: 1.0e-3)

Returns

The state at time t1.

solve_ltc(x, activation, tau, opts \\ [])

@spec solve_ltc(Nx.Tensor.t(), Nx.Tensor.t(), Nx.Tensor.t(), keyword()) ::
  Nx.Tensor.t()

Solve the LTC (Liquid Time-Constant) ODE for one timestep.

The LTC ODE is: dx/dt = (-x + activation) / tau

This is optimized for Liquid Neural Networks where we integrate the state for a single frame (dt = 1.0 normalized).

Parameters

x - Current hidden state [batch, hidden_size]
activation - The f(x, input) activation [batch, hidden_size]
tau - Time constants [batch, hidden_size]
opts - Solver options

Options

:solver - Solver type (default: :rk4)
:steps - Number of integration sub-steps (default: 1)

Returns

The new hidden state after one frame.