Skip to content

Optimizers

13 built-in optimizers, dispatched into 5 families at construction time (before JIT). Each family owns its own grad-step factory in optimizers/<family>.py; the generic step is make_grad_step_standard.

Family Names Step shape
STANDARD adam, sgd, adamw, lion, muon, ngd, shampoo jax.grad → opt.update(grads, state, params)
PCGRAD (gradient_surgery: pcgrad) Per-equation grads with conflict projection
MAO mao, mao_kfac Per-equation Jacobian via jax.jacrev → MAO update
LBFGS lbfgs Optax LBFGS with line search (needs value, grad, value_fn)
GN gn, ign, lm Gauss-Newton / Levenberg-Marquardt: Δθ = −(JᵀJ)⁻¹ Jᵀr

Registration uses the @register_optimizer(name, kind) decorator in each module; optimizers/__init__.py imports every module to trigger registration. create_optimizer(config) looks up the registry and chains optax.clip_by_global_norm for STANDARD optimizers automatically when grad_clip is set.

MAO uses _MAOFactory for deferred n_tasks resolution (the model's equation count is known only at create_train_state time).

Composite loss is rejected with MAO/GN/IGN/LM/LBFGS and PCGrad (TrainConfig._validate_ranges enforces this); on those paths the optimizer's update doesn't see the auxiliary terms.

For adding a new optimizer, see Adding an optimizer.

deqn_jax.optimizers.registry

Optimizer registry for DEQN-JAX.

Maps optimizer names to factory functions and optimizer kinds. The kind determines which train_step variant is used.

OptimizerKind

Bases: str, Enum

Determines which train_step variant runs.

Source code in src/deqn_jax/optimizers/registry.py 
13
14
15
16
17
18
19
class OptimizerKind(str, Enum):
    """Determines which train_step variant runs."""

    STANDARD = "standard"  # adam, sgd, adamw, lion, muon, ngd, shampoo, kfac
    MAO = "mao"  # per-equation Jacobian
    LBFGS = "lbfgs"  # extra args for line search
    GN = "gn"  # Gauss-Newton / LM (needs residual_fn)

ReduceLROnPlateau

Keras-style loss-reactive LR decay.

Mirrors DEQN-MAO's lr_scheduler: ReduceLROnPlateau with the same knob names (factor, patience, cooldown, min_delta, min_lr). Unlike the other schedules here, this one is stateful and must be called with the current loss at every cycle, not just the step index.

Call signature: self(ep_num, loss=None) -> lr. loss=None is treated as "no update this cycle" -- the scheduler returns the current LR without advancing state. Used so the schedule callable can be probed early in training (e.g. for logging) before a first loss is available.

Source code in src/deqn_jax/optimizers/registry.py 
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
class ReduceLROnPlateau:
    """Keras-style loss-reactive LR decay.

    Mirrors DEQN-MAO's ``lr_scheduler: ReduceLROnPlateau`` with the same
    knob names (factor, patience, cooldown, min_delta, min_lr). Unlike
    the other schedules here, this one is stateful and must be called
    with the current loss at every cycle, not just the step index.

    Call signature: ``self(ep_num, loss=None) -> lr``. ``loss=None`` is
    treated as "no update this cycle" -- the scheduler returns the
    current LR without advancing state. Used so the schedule callable
    can be probed early in training (e.g. for logging) before a first
    loss is available.
    """

    def __init__(self, initial_lr, factor, patience, cooldown, min_delta, min_lr):
        self.initial_lr = float(initial_lr)
        self.factor = float(factor)
        self.patience = int(patience)
        self.cooldown = int(cooldown)
        self.min_delta = float(min_delta)
        self.min_lr = float(min_lr)

        self._lr = self.initial_lr
        self._best = float("inf")
        self._wait = 0
        self._cooldown_counter = 0

    def __call__(self, _ep_num=None, loss=None):
        if loss is None:
            return self._lr

        # In cooldown: don't change LR, just tick down.
        if self._cooldown_counter > 0:
            self._cooldown_counter -= 1
            self._best = min(self._best, loss)
            self._wait = 0
            return self._lr

        # Track improvement.
        if loss < self._best - self.min_delta:
            self._best = loss
            self._wait = 0
        else:
            self._wait += 1
            if self._wait >= self.patience:
                new_lr = max(self._lr * self.factor, self.min_lr)
                self._lr = new_lr
                self._cooldown_counter = self.cooldown
                self._wait = 0
        return self._lr

register_optimizer 

register_optimizer(
    name: str, kind: OptimizerKind = OptimizerKind.STANDARD
)

Decorator to register an optimizer factory.

Source code in src/deqn_jax/optimizers/registry.py 
26
27
28
29
30
31
32
33
34
35
36
def register_optimizer(
    name: str,
    kind: OptimizerKind = OptimizerKind.STANDARD,
):
    """Decorator to register an optimizer factory."""

    def decorator(fn: Callable) -> Callable:
        _REGISTRY[name] = (fn, kind)
        return fn

    return decorator

create_optimizer 

create_optimizer(
    config, total_steps: Optional[int] = None
) -> Tuple[Any, OptimizerKind]

Create optimizer from config.

LR schedules are NOT applied here — they break XLA kernel fusion and cause 5-6x slowdowns. Instead, the training loop periodically recreates the optimizer with an updated constant LR. See _build_lr_schedule for computing schedule values and train_from_config for the periodic update logic.

Parameters:

Name Type Description Default
config

OptimizerConfig with at least a name field.

 required
total_steps Optional[int]

Unused (kept for API compat). Schedule is handled by the training loop.

 None

Returns:

Type Description
Tuple[Any, OptimizerKind]

Tuple of (optimizer, OptimizerKind)
Source code in src/deqn_jax/optimizers/registry.py 
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
def create_optimizer(
    config,
    total_steps: Optional[int] = None,
) -> Tuple[Any, OptimizerKind]:
    """Create optimizer from config.

    LR schedules are NOT applied here — they break XLA kernel fusion
    and cause 5-6x slowdowns. Instead, the training loop periodically
    recreates the optimizer with an updated constant LR. See
    ``_build_lr_schedule`` for computing schedule values and
    ``train_from_config`` for the periodic update logic.

    Args:
        config: OptimizerConfig with at least a ``name`` field.
        total_steps: Unused (kept for API compat). Schedule is handled
            by the training loop.

    Returns:
        Tuple of (optimizer, OptimizerKind)
    """
    name = config.name
    if name not in _REGISTRY:
        available = ", ".join(sorted(_REGISTRY.keys()))
        raise ValueError(f"Unknown optimizer '{name}'. Available: {available}")

    factory, kind = _REGISTRY[name]
    opt = factory(config)

    # Chain grad clip for STANDARD optimizers
    if kind == OptimizerKind.STANDARD and config.grad_clip is not None:
        opt = optax.chain(optax.clip_by_global_norm(config.grad_clip), opt)

    return opt, kind

list_optimizers 

list_optimizers() -> List[str]

Return sorted list of registered optimizer names.

Source code in src/deqn_jax/optimizers/registry.py 
176
177
178
def list_optimizers() -> List[str]:
    """Return sorted list of registered optimizer names."""
    return sorted(_REGISTRY.keys())

deqn_jax.optimizers.ngd

Natural Gradient Descent (diagonal Fisher approximation).

Running diagonal Fisher via EMA of g², preconditioned step: θ ← θ - lr * g / (sqrt(F) + damping)

Cheap and effective for PINN-style losses.

NGDState

Bases: NamedTuple

State for diagonal Fisher NGD.

Source code in src/deqn_jax/optimizers/ngd.py 
19
20
21
22
23
class NGDState(NamedTuple):
    """State for diagonal Fisher NGD."""

    count: Array
    fisher_diag: Any  # pytree matching params, EMA of g²

ngd 

ngd(
    learning_rate: float = 0.001,
    damping: float = 0.0001,
    decay: float = 0.999,
) -> optax.GradientTransformation

Diagonal Fisher Natural Gradient Descent.

Parameters:

Name Type Description Default
learning_rate float

Step size

 0.001
damping float

Regularization added to sqrt(Fisher)

 0.0001
decay float

EMA decay for Fisher diagonal estimate

 0.999

Returns:

Type Description
GradientTransformation

optax.GradientTransformation
Source code in src/deqn_jax/optimizers/ngd.py 
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
def ngd(
    learning_rate: float = 1e-3,
    damping: float = 1e-4,
    decay: float = 0.999,
) -> optax.GradientTransformation:
    """Diagonal Fisher Natural Gradient Descent.

    Args:
        learning_rate: Step size
        damping: Regularization added to sqrt(Fisher)
        decay: EMA decay for Fisher diagonal estimate

    Returns:
        optax.GradientTransformation
    """

    def init_fn(params) -> NGDState:
        fisher_diag = jax.tree.map(jnp.zeros_like, params)
        return NGDState(count=jnp.zeros([], dtype=jnp.int32), fisher_diag=fisher_diag)

    def update_fn(
        updates: Any,
        state: NGDState,
        params: Optional[Any] = None,
    ) -> Tuple[Any, NGDState]:
        # Update Fisher diagonal: F ← decay * F + (1-decay) * g²
        new_fisher = jax.tree.map(
            lambda f, g: decay * f + (1.0 - decay) * g**2,
            state.fisher_diag,
            updates,
        )
        # Preconditioned step: -lr * g / (sqrt(F) + damping)
        preconditioned = jax.tree.map(
            lambda g, f: -learning_rate * g / (jnp.sqrt(f) + damping),
            updates,
            new_fisher,
        )
        return preconditioned, NGDState(count=state.count + 1, fisher_diag=new_fisher)

    return optax.GradientTransformation(init_fn, update_fn)

deqn_jax.optimizers.mao

Multi-Adaptive Optimizer (MAO) for per-equation optimization.

MAO maintains separate Adam-style moment estimates for each equation, then combines updates via per-task adaptive learning rates.

This is NOT an optax.GradientTransformation -- it has a custom interface because it receives per-equation Jacobians instead of standard gradients.

Usage in training

eq_jac = jax.jacrev(per_eq_loss_fn)(params) # pytree, each leaf [n_eq, *shape] updates, new_state = mao.update(eq_jac, state, params) new_params = optax.apply_updates(params, updates)

MAOState

Bases: NamedTuple

State for MAO optimizer.

Source code in src/deqn_jax/optimizers/mao.py 
24
25
26
27
28
29
class MAOState(NamedTuple):
    """State for MAO optimizer."""

    count: Array  # scalar step count
    m: Any  # first moment pytree, each leaf [n_eq, *param_shape]
    v: Any  # second moment pytree, each leaf [n_eq, *param_shape]

MAOTransform

Multi-Adaptive Optimizer.

Maintains per-equation moment estimates and combines updates.

Source code in src/deqn_jax/optimizers/mao.py 
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
class MAOTransform:
    """Multi-Adaptive Optimizer.

    Maintains per-equation moment estimates and combines updates.
    """

    def __init__(
        self,
        learning_rate: float = 1e-3,
        beta1: float = 0.9,
        beta2: float = 0.999,
        epsilon: float = 1e-8,
        n_tasks: int = 1,
    ):
        self.learning_rate = learning_rate
        self.beta1 = beta1
        self.beta2 = beta2
        self.epsilon = epsilon
        self.n_tasks = n_tasks

    def init(self, params: Any) -> MAOState:
        """Initialize MAO state with per-equation moments."""
        m = jax.tree.map(
            lambda p: jnp.zeros((self.n_tasks,) + p.shape, dtype=p.dtype),
            params,
        )
        v = jax.tree.map(
            lambda p: jnp.zeros((self.n_tasks,) + p.shape, dtype=p.dtype),
            params,
        )
        return MAOState(count=jnp.zeros([], dtype=jnp.int32), m=m, v=v)

    def update(
        self,
        eq_jacobian: Any,
        state: MAOState,
        params: Any,
    ) -> Tuple[Any, MAOState]:
        """Compute MAO update from per-equation Jacobians.

        Args:
            eq_jacobian: Pytree matching params, each leaf [n_eq, *param_shape]
                        (output of jax.jacrev over per-equation losses)
            state: Current MAO state
            params: Current parameters (unused but kept for API consistency)

        Returns:
            Tuple of (updates pytree, new_state)
        """
        count = state.count + 1
        b1, b2, eps = self.beta1, self.beta2, self.epsilon

        # Update per-equation moments
        new_m = jax.tree.map(
            lambda m, j: b1 * m + (1.0 - b1) * j,
            state.m,
            eq_jacobian,
        )
        new_v = jax.tree.map(
            lambda v, j: b2 * v + (1.0 - b2) * j**2,
            state.v,
            eq_jacobian,
        )

        # Bias correction
        bc1 = 1.0 - b1**count
        bc2 = 1.0 - b2**count

        # Per-equation Adam updates, then sum across equations
        def compute_update(m_leaf, v_leaf):
            # m_leaf: [n_eq, *shape], v_leaf: [n_eq, *shape]
            m_hat = m_leaf / bc1
            v_hat = v_leaf / bc2
            # Per-equation update: [n_eq, *shape]
            per_eq = m_hat / (jnp.sqrt(v_hat) + eps)
            # Average across equations → [*shape]
            return -self.learning_rate * jnp.mean(per_eq, axis=0)

        updates = jax.tree.map(compute_update, new_m, new_v)

        new_state = MAOState(count=count, m=new_m, v=new_v)
        return updates, new_state

init 

init(params: Any) -> MAOState

Initialize MAO state with per-equation moments.

Source code in src/deqn_jax/optimizers/mao.py 
52
53
54
55
56
57
58
59
60
61
62
def init(self, params: Any) -> MAOState:
    """Initialize MAO state with per-equation moments."""
    m = jax.tree.map(
        lambda p: jnp.zeros((self.n_tasks,) + p.shape, dtype=p.dtype),
        params,
    )
    v = jax.tree.map(
        lambda p: jnp.zeros((self.n_tasks,) + p.shape, dtype=p.dtype),
        params,
    )
    return MAOState(count=jnp.zeros([], dtype=jnp.int32), m=m, v=v)

update 

update(
    eq_jacobian: Any, state: MAOState, params: Any
) -> Tuple[Any, MAOState]

Compute MAO update from per-equation Jacobians.

Parameters:

Name Type Description Default
eq_jacobian Any

Pytree matching params, each leaf [n_eq, *param_shape] (output of jax.jacrev over per-equation losses)

 required
state MAOState

Current MAO state

 required
params Any

Current parameters (unused but kept for API consistency)

 required

Returns:

Type Description
Tuple[Any, MAOState]

Tuple of (updates pytree, new_state)
Source code in src/deqn_jax/optimizers/mao.py 
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
def update(
    self,
    eq_jacobian: Any,
    state: MAOState,
    params: Any,
) -> Tuple[Any, MAOState]:
    """Compute MAO update from per-equation Jacobians.

    Args:
        eq_jacobian: Pytree matching params, each leaf [n_eq, *param_shape]
                    (output of jax.jacrev over per-equation losses)
        state: Current MAO state
        params: Current parameters (unused but kept for API consistency)

    Returns:
        Tuple of (updates pytree, new_state)
    """
    count = state.count + 1
    b1, b2, eps = self.beta1, self.beta2, self.epsilon

    # Update per-equation moments
    new_m = jax.tree.map(
        lambda m, j: b1 * m + (1.0 - b1) * j,
        state.m,
        eq_jacobian,
    )
    new_v = jax.tree.map(
        lambda v, j: b2 * v + (1.0 - b2) * j**2,
        state.v,
        eq_jacobian,
    )

    # Bias correction
    bc1 = 1.0 - b1**count
    bc2 = 1.0 - b2**count

    # Per-equation Adam updates, then sum across equations
    def compute_update(m_leaf, v_leaf):
        # m_leaf: [n_eq, *shape], v_leaf: [n_eq, *shape]
        m_hat = m_leaf / bc1
        v_hat = v_leaf / bc2
        # Per-equation update: [n_eq, *shape]
        per_eq = m_hat / (jnp.sqrt(v_hat) + eps)
        # Average across equations → [*shape]
        return -self.learning_rate * jnp.mean(per_eq, axis=0)

    updates = jax.tree.map(compute_update, new_m, new_v)

    new_state = MAOState(count=count, m=new_m, v=new_v)
    return updates, new_state

make_grad_step_mao 

make_grad_step_mao(
    model,
    mao_opt: Any,
    mc_samples: int,
    quad_nodes,
    quad_weights,
    loss_reweight: str,
    reweight_alpha: float,
    use_target_network: bool,
    compute_loss_fn,
    grad_clip,
)

JIT'd: one MAO (per-equation Jacobian) gradient update on a minibatch.

Source code in src/deqn_jax/optimizers/mao.py 
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
def make_grad_step_mao(
    model,
    mao_opt: Any,
    mc_samples: int,
    quad_nodes,
    quad_weights,
    loss_reweight: str,
    reweight_alpha: float,
    use_target_network: bool,
    compute_loss_fn,
    grad_clip,
):
    """JIT'd: one MAO (per-equation Jacobian) gradient update on a minibatch."""
    import equinox as eqx
    import optax

    from deqn_jax.training.loss import compute_loss, eq_losses_to_array
    from deqn_jax.training.reweighting import update_reweighting
    from deqn_jax.types import Metrics, TrainState

    n_eq = len(model.equation_names) if model.equation_names else 1
    _compute_loss_total = compute_loss_fn or compute_loss

    @jax.jit
    def grad_step(
        state: TrainState,
        batch: Array,
        lr_scale: Array,
        shock_scale: Array = jnp.array(1.0),
    ) -> Tuple[TrainState, Metrics]:
        loss_key, new_key = jax.random.split(state.key)
        target_fn = state.target_params if use_target_network else None

        params_arrays = eqx.filter(state.params, eqx.is_array)
        params_static = eqx.filter(state.params, lambda x: not eqx.is_array(x))

        def per_eq_loss_fn(p_arrays):
            full_params = eqx.combine(p_arrays, params_static)
            _, eq_losses = compute_loss(
                model,
                full_params,
                batch,
                loss_key,
                mc_samples,
                weights=None,
                shock_scale=shock_scale,
                quad_nodes=quad_nodes,
                quad_weights=quad_weights,
                target_policy_fn=target_fn,
            )
            return eq_losses_to_array(eq_losses)

        eq_jac = jax.jacrev(per_eq_loss_fn)(params_arrays)

        def total_loss_fn(params):
            loss, eq_losses = _compute_loss_total(
                model,
                params,
                batch,
                loss_key,
                mc_samples,
                weights=state.loss_weights,
                shock_scale=shock_scale,
                quad_nodes=quad_nodes,
                quad_weights=quad_weights,
                target_policy_fn=target_fn,
            )
            return loss, eq_losses

        (loss, eq_losses), grads = eqx.filter_value_and_grad(
            total_loss_fn, has_aux=True
        )(state.params)
        grad_norm = optax.global_norm(eqx.filter(grads, eqx.is_array))

        updates, new_opt_state = mao_opt.update(eq_jac, state.opt_state, params_arrays)
        if grad_clip is not None:
            update_norm = optax.global_norm(updates)
            clip_scale = jnp.minimum(1.0, grad_clip / (update_norm + 1e-8))
            updates = jax.tree.map(lambda u: clip_scale * u, updates)
        updates = jax.tree.map(lambda u: lr_scale * u, updates)
        new_params_arrays = optax.apply_updates(params_arrays, updates)
        new_params = eqx.combine(new_params_arrays, state.params)

        new_weights, new_rw = update_reweighting(
            eq_losses,
            state,
            loss_reweight,
            reweight_alpha,
            n_eq,
        )
        new_state = state._replace(
            params=new_params,
            opt_state=new_opt_state,
            key=new_key,
            step=state.step + 1,
            loss_weights=new_weights,
            reweight_state=new_rw,
        )
        return new_state, Metrics(loss=loss, residuals=eq_losses, grad_norm=grad_norm)

    return grad_step

deqn_jax.optimizers.shampoo

In-house Kronecker-factored Shampoo optimizer.

For 2D parameters (weight matrices), maintains Kronecker factors: L = EMA(G @ G^T), R = EMA(G^T @ G) update = L^{-1/4} @ G @ R^{-1/4}

For 1D parameters (biases), reshapes to [1, n] to avoid data-dependent branching.

Preconditioners are updated every precond_update_freq steps.

ShampooState

Bases: NamedTuple

Shampoo optimizer state.

Source code in src/deqn_jax/optimizers/shampoo.py 
22
23
24
25
26
27
class ShampooState(NamedTuple):
    """Shampoo optimizer state."""

    count: Array
    L: Any  # pytree of left preconditioners
    R: Any  # pytree of right preconditioners

shampoo 

shampoo(
    learning_rate: float = 0.001,
    beta: float = 0.9,
    precond_update_freq: int = 10,
    epsilon: float = 1e-12,
) -> optax.GradientTransformation

Kronecker-factored Shampoo optimizer.

Parameters:

Name Type Description Default
learning_rate float

Step size

 0.001
beta float

EMA decay for preconditioner statistics

 0.9
precond_update_freq int

Steps between preconditioner updates

 10
epsilon float

Ridge for numerical stability

 1e-12

Returns:

Type Description
GradientTransformation

optax.GradientTransformation
Source code in src/deqn_jax/optimizers/shampoo.py 
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
def shampoo(
    learning_rate: float = 1e-3,
    beta: float = 0.9,
    precond_update_freq: int = 10,
    epsilon: float = 1e-12,
) -> optax.GradientTransformation:
    """Kronecker-factored Shampoo optimizer.

    Args:
        learning_rate: Step size
        beta: EMA decay for preconditioner statistics
        precond_update_freq: Steps between preconditioner updates
        epsilon: Ridge for numerical stability

    Returns:
        optax.GradientTransformation
    """

    def init_fn(params) -> ShampooState:
        def make_L(p):
            if p.ndim < 2:
                return jnp.eye(1, dtype=p.dtype)
            return jnp.eye(p.shape[0], dtype=p.dtype)

        def make_R(p):
            if p.ndim < 2:
                n = p.shape[0] if p.ndim == 1 else 1
                return jnp.eye(n, dtype=p.dtype)
            return jnp.eye(p.shape[1], dtype=p.dtype)

        L = jax.tree.map(make_L, params)
        R = jax.tree.map(make_R, params)
        return ShampooState(count=jnp.zeros([], dtype=jnp.int32), L=L, R=R)

    def update_fn(
        updates: Any,
        state: ShampooState,
        params: Optional[Any] = None,
    ) -> Tuple[Any, ShampooState]:
        count = state.count + 1
        do_update = (count % precond_update_freq) == 0

        def update_L(g, L_old):
            if g.ndim < 2:
                g_2d = g.reshape(1, -1)
            else:
                g_2d = g
            return jax.lax.cond(
                do_update,
                lambda _: beta * L_old + (1.0 - beta) * (g_2d @ g_2d.T),
                lambda _: L_old,
                None,
            )

        def update_R(g, R_old):
            if g.ndim < 2:
                g_2d = g.reshape(1, -1)
            else:
                g_2d = g
            return jax.lax.cond(
                do_update,
                lambda _: beta * R_old + (1.0 - beta) * (g_2d.T @ g_2d),
                lambda _: R_old,
                None,
            )

        def precondition(g, L_new, R_new):
            original_shape = g.shape
            if g.ndim < 2:
                g_2d = g.reshape(1, -1)
            else:
                g_2d = g

            L_inv4 = _matrix_power_neg_quarter(L_new, ridge=epsilon)
            R_inv4 = _matrix_power_neg_quarter(R_new, ridge=epsilon)
            precond = L_inv4 @ g_2d @ R_inv4

            if g.ndim < 2:
                return -learning_rate * precond.reshape(original_shape)
            return -learning_rate * precond

        new_L = jax.tree.map(update_L, updates, state.L)
        new_R = jax.tree.map(update_R, updates, state.R)
        new_updates = jax.tree.map(precondition, updates, new_L, new_R)

        return new_updates, ShampooState(count=count, L=new_L, R=new_R)

    return optax.GradientTransformation(init_fn, update_fn)

deqn_jax.optimizers.lbfgs

L-BFGS optimizer via optax.

Thin wrapper around optax.lbfgs() which is a GradientTransformationExtraArgs -- it needs value and value_fn passed to update() for line search.

make_grad_step_lbfgs 

make_grad_step_lbfgs(
    model: ModelSpec,
    opt: Any,
    mc_samples: int,
    quad_nodes: Optional[Array],
    quad_weights: Optional[Array],
    loss_reweight: str,
    reweight_alpha: float,
    use_target_network: bool,
    compute_loss_fn: Optional[Callable],
)

JIT'd: one L-BFGS gradient update on a minibatch (with line search).

Source code in src/deqn_jax/optimizers/lbfgs.py 
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
def make_grad_step_lbfgs(
    model: ModelSpec,
    opt: Any,
    mc_samples: int,
    quad_nodes: Optional[Array],
    quad_weights: Optional[Array],
    loss_reweight: str,
    reweight_alpha: float,
    use_target_network: bool,
    compute_loss_fn: Optional[Callable],
):
    """JIT'd: one L-BFGS gradient update on a minibatch (with line search)."""
    n_eq = len(model.equation_names) if model.equation_names else 1
    _compute_loss = compute_loss_fn or compute_loss

    @jax.jit
    def grad_step(
        state: TrainState,
        batch: Array,
        lr_scale: Array,
        shock_scale: Array = jnp.array(1.0),
    ) -> Tuple[TrainState, Metrics]:
        loss_key, new_key = jax.random.split(state.key)
        target_fn = state.target_params if use_target_network else None

        params_arrays = eqx.filter(state.params, eqx.is_array)
        params_static = eqx.filter(state.params, lambda x: not eqx.is_array(x))

        def loss_fn(params):
            loss, eq_losses = _compute_loss(
                model,
                params,
                batch,
                loss_key,
                mc_samples,
                weights=state.loss_weights,
                shock_scale=shock_scale,
                quad_nodes=quad_nodes,
                quad_weights=quad_weights,
                target_policy_fn=target_fn,
            )
            return loss, eq_losses

        (loss, eq_losses), grads = eqx.filter_value_and_grad(loss_fn, has_aux=True)(
            state.params
        )
        grads_arrays = eqx.filter(grads, eqx.is_array)
        grad_norm = optax.global_norm(grads_arrays)

        def value_fn(p_arrays):
            full_params = eqx.combine(p_arrays, params_static)
            v, _ = _compute_loss(
                model,
                full_params,
                batch,
                loss_key,
                mc_samples,
                weights=state.loss_weights,
                shock_scale=shock_scale,
                quad_nodes=quad_nodes,
                quad_weights=quad_weights,
                target_policy_fn=target_fn,
            )
            return v

        updates, new_opt_state = opt.update(
            grads_arrays,
            state.opt_state,
            params_arrays,
            value=loss,
            grad=grads_arrays,
            value_fn=value_fn,
        )
        updates = jax.tree.map(lambda u: lr_scale * u, updates)
        new_params_arrays = optax.apply_updates(params_arrays, updates)
        new_params = eqx.combine(new_params_arrays, state.params)

        new_weights, new_rw = update_reweighting(
            eq_losses,
            state,
            loss_reweight,
            reweight_alpha,
            n_eq,
        )
        new_state = state._replace(
            params=new_params,
            opt_state=new_opt_state,
            key=new_key,
            step=state.step + 1,
            loss_weights=new_weights,
            reweight_state=new_rw,
        )
        return new_state, Metrics(loss=loss, residuals=eq_losses, grad_norm=grad_norm)

    return grad_step

deqn_jax.optimizers.gauss_newton

Gauss-Newton and Levenberg-Marquardt optimizers for JAX.

For DEQN, we minimize L = ||r(θ)||² where r are equilibrium residuals. Gauss-Newton approximates Hessian as H ≈ J^T J where J = ∂r/∂θ.

Key advantage over first-order methods: quadratic convergence near solution.

This implementation uses JAX autodiff (jacrev/jacfwd) for efficient Jacobian computation - much faster than finite differences.

Note on optimistix: optimistix's LevenbergMarquardt and GaussNewton solvers are excellent for full-convergence least-squares problems (curve fitting, SS solving) and are used in this codebase for those cases (see models/disaster/equations.py:solve_omega_bar and models/disaster/steady_state.py). They do not fit DEQN's per-minibatch optimizer contract, where each update() is called once with a fresh residual function (different minibatch's residuals) and we want exactly one step. Optimistix's trust-region update is conservative on first call (init→step counts as ≥1 step internally, returning y0 with max_steps=1; max_steps=2 takes only a half-step) and re-initializing per call defeats the purpose. So the classes below remain hand-rolled.

Usage

opt = gauss_newton(learning_rate=1.0) state = opt.init(params)

def residual_fn(p): return model.equations(states, p(states), ...)

params, state = opt.update(residual_fn, params, state)

GaussNewtonState

Bases: NamedTuple

State for Gauss-Newton optimizer.

last_loss is a JAX scalar Array at runtime (sum of squared residuals from inside the JIT'd update step). It was annotated as Python float originally, which produced spurious invalid-argument-type errors at every constructor call — same pattern as the Metrics annotation lie cleared in commit 3ae741f. damping stays float because the LM update keeps it on the Python side.

Source code in src/deqn_jax/optimizers/gauss_newton.py 
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
class GaussNewtonState(NamedTuple):
    """State for Gauss-Newton optimizer.

    ``last_loss`` is a JAX scalar Array at runtime (sum of squared
    residuals from inside the JIT'd update step). It was annotated as
    Python ``float`` originally, which produced spurious
    ``invalid-argument-type`` errors at every constructor call — same
    pattern as the ``Metrics`` annotation lie cleared in commit
    ``3ae741f``. ``damping`` stays ``float`` because the LM update
    keeps it on the Python side.
    """

    count: int  # Iteration count
    damping: float  # Current LM damping
    last_loss: Array  # Previous loss for adaptive damping

ImplicitGaussNewtonState

Bases: NamedTuple

State for matrix-free damped Gauss-Newton.

last_cg_residual tracks the final linear-system residual norm from the conjugate-gradient solve. It is diagnostic only; schedules and checkpoint serialization treat it like the other scalar optimizer-state arrays.

Source code in src/deqn_jax/optimizers/gauss_newton.py 
59
60
61
62
63
64
65
66
67
68
69
70
71
class ImplicitGaussNewtonState(NamedTuple):
    """State for matrix-free damped Gauss-Newton.

    ``last_cg_residual`` tracks the final linear-system residual norm from the
    conjugate-gradient solve. It is diagnostic only; schedules and checkpoint
    serialization treat it like the other scalar optimizer-state arrays.
    """

    count: int
    damping: float
    last_loss: Array
    last_cg_residual: Array
    last_cg_iters: Array

GaussNewton

Gauss-Newton optimizer for nonlinear least squares.

Source code in src/deqn_jax/optimizers/gauss_newton.py 
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
class GaussNewton:
    """Gauss-Newton optimizer for nonlinear least squares."""

    def __init__(
        self,
        learning_rate: float = 1.0,
        damping: float = 0.0,
    ):
        self.learning_rate = learning_rate
        self.damping = damping

    def init(self, params) -> GaussNewtonState:
        # ``count`` must be a JAX scalar from the start; after the first
        # update it becomes one anyway (jnp ops upcast Python ints), and
        # checkpoint round-trip needs the *initial* and *post-update* leaf
        # types to match the saved tree.
        return GaussNewtonState(
            count=jnp.asarray(0, dtype=jnp.int32),
            damping=self.damping,
            last_loss=jnp.asarray(jnp.inf),
        )

    def update(
        self,
        residual_fn: Callable,
        params: Any,
        state: GaussNewtonState,
        lr_scale: Any = 1.0,
    ) -> Tuple[Any, GaussNewtonState]:
        """Perform one GN step.

        Args:
            residual_fn: Function params -> flat residuals [n_residuals]
            params: Current parameters (pytree)
            state: Optimizer state

        Returns:
            Tuple of (new_params, new_state)
        """
        # Flatten params for linear algebra
        flat_params, unflatten = jax.flatten_util.ravel_pytree(params)
        n_params = flat_params.shape[0]

        # Wrap residual_fn to work with flat params
        def flat_residual_fn(flat_p):
            p = unflatten(flat_p)
            r = residual_fn(p)
            return jnp.ravel(r)

        # Compute residuals
        r = flat_residual_fn(flat_params)
        n_residuals = r.shape[0]

        # Compute Jacobian using autodiff
        # Choose forward or reverse mode based on dimensions
        if n_residuals <= n_params:
            J = jax.jacrev(flat_residual_fn)(flat_params)
        else:
            J = jax.jacfwd(flat_residual_fn)(flat_params)

        # Ensure minimum damping for numerical stability
        damping = jnp.maximum(state.damping, 1e-6)

        # Solve (J^T J + λI) δ = -J^T r
        # When n_residuals < n_params, use dual formulation (Woodbury):
        #   δ = -J^T (J J^T + λI)^{-1} r
        # Solves (n_res, n_res) system instead of (n_params, n_params).
        if n_residuals < n_params:
            G = J @ J.T + damping * jnp.eye(n_residuals)
            v = jnp.linalg.solve(G, r)
            delta = -J.T @ v
        else:
            JtJ = J.T @ J + damping * jnp.eye(n_params)
            delta = jnp.linalg.solve(JtJ, -(J.T @ r))

        # Match the train_step contract used by the other optimizers: when a
        # schedule is active, self.learning_rate is 1.0 and lr_scale carries
        # the per-step learning rate.
        step_size = self.learning_rate * lr_scale
        new_flat_params = flat_params + step_size * delta
        new_params = unflatten(new_flat_params)

        # Compute new loss (from updated params, not old residuals)
        new_r = flat_residual_fn(new_flat_params)
        new_loss = jnp.sum(new_r**2)

        new_state = GaussNewtonState(
            count=state.count + 1,
            damping=state.damping,
            last_loss=new_loss,
        )

        return new_params, new_state

update 

update(
    residual_fn: Callable,
    params: Any,
    state: GaussNewtonState,
    lr_scale: Any = 1.0,
) -> Tuple[Any, GaussNewtonState]

Perform one GN step.

Parameters:

Name Type Description Default
residual_fn Callable

Function params -> flat residuals [n_residuals]

 required
params Any

Current parameters (pytree)

 required
state GaussNewtonState

Optimizer state

 required

Returns:

Type Description
Tuple[Any, GaussNewtonState]

Tuple of (new_params, new_state)
Source code in src/deqn_jax/optimizers/gauss_newton.py 
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
def update(
    self,
    residual_fn: Callable,
    params: Any,
    state: GaussNewtonState,
    lr_scale: Any = 1.0,
) -> Tuple[Any, GaussNewtonState]:
    """Perform one GN step.

    Args:
        residual_fn: Function params -> flat residuals [n_residuals]
        params: Current parameters (pytree)
        state: Optimizer state

    Returns:
        Tuple of (new_params, new_state)
    """
    # Flatten params for linear algebra
    flat_params, unflatten = jax.flatten_util.ravel_pytree(params)
    n_params = flat_params.shape[0]

    # Wrap residual_fn to work with flat params
    def flat_residual_fn(flat_p):
        p = unflatten(flat_p)
        r = residual_fn(p)
        return jnp.ravel(r)

    # Compute residuals
    r = flat_residual_fn(flat_params)
    n_residuals = r.shape[0]

    # Compute Jacobian using autodiff
    # Choose forward or reverse mode based on dimensions
    if n_residuals <= n_params:
        J = jax.jacrev(flat_residual_fn)(flat_params)
    else:
        J = jax.jacfwd(flat_residual_fn)(flat_params)

    # Ensure minimum damping for numerical stability
    damping = jnp.maximum(state.damping, 1e-6)

    # Solve (J^T J + λI) δ = -J^T r
    # When n_residuals < n_params, use dual formulation (Woodbury):
    #   δ = -J^T (J J^T + λI)^{-1} r
    # Solves (n_res, n_res) system instead of (n_params, n_params).
    if n_residuals < n_params:
        G = J @ J.T + damping * jnp.eye(n_residuals)
        v = jnp.linalg.solve(G, r)
        delta = -J.T @ v
    else:
        JtJ = J.T @ J + damping * jnp.eye(n_params)
        delta = jnp.linalg.solve(JtJ, -(J.T @ r))

    # Match the train_step contract used by the other optimizers: when a
    # schedule is active, self.learning_rate is 1.0 and lr_scale carries
    # the per-step learning rate.
    step_size = self.learning_rate * lr_scale
    new_flat_params = flat_params + step_size * delta
    new_params = unflatten(new_flat_params)

    # Compute new loss (from updated params, not old residuals)
    new_r = flat_residual_fn(new_flat_params)
    new_loss = jnp.sum(new_r**2)

    new_state = GaussNewtonState(
        count=state.count + 1,
        damping=state.damping,
        last_loss=new_loss,
    )

    return new_params, new_state

ImplicitGaussNewton

Matrix-free damped Gauss-Newton / natural-gradient optimizer.

For residual least squares 0.5 * ||r(theta)||^2, the Gauss-Newton metric is J.T @ J where J = dr/dtheta. This class solves

``(J.T @ J + damping * I) delta = -J.T @ r``

with conjugate gradients using only JVP/VJP products. It avoids storing the dense Fisher/GN matrix and avoids materializing the full residual Jacobian, which is the practical route for DEQN-sized policy networks.

Source code in src/deqn_jax/optimizers/gauss_newton.py 
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
class ImplicitGaussNewton:
    """Matrix-free damped Gauss-Newton / natural-gradient optimizer.

    For residual least squares ``0.5 * ||r(theta)||^2``, the Gauss-Newton
    metric is ``J.T @ J`` where ``J = dr/dtheta``. This class solves

        ``(J.T @ J + damping * I) delta = -J.T @ r``

    with conjugate gradients using only JVP/VJP products. It avoids storing the
    dense Fisher/GN matrix and avoids materializing the full residual Jacobian,
    which is the practical route for DEQN-sized policy networks.
    """

    def __init__(
        self,
        learning_rate: float = 1.0,
        damping: float = 1e-4,
        cg_iters: int = 20,
        cg_tol: float = 1e-6,
    ):
        self.learning_rate = learning_rate
        self.damping = damping
        self.cg_iters = int(cg_iters)
        self.cg_tol = float(cg_tol)

    def init(self, params) -> ImplicitGaussNewtonState:
        return ImplicitGaussNewtonState(
            count=jnp.asarray(0, dtype=jnp.int32),
            damping=self.damping,
            last_loss=jnp.asarray(jnp.inf),
            last_cg_residual=jnp.asarray(jnp.inf),
            last_cg_iters=jnp.asarray(0, dtype=jnp.int32),
        )

    def update(
        self,
        residual_fn: Callable,
        params: Any,
        state: ImplicitGaussNewtonState,
        lr_scale: Any = 1.0,
    ) -> Tuple[Any, ImplicitGaussNewtonState]:
        flat_params, unflatten = jax.flatten_util.ravel_pytree(params)

        def flat_residual_fn(flat_p):
            p = unflatten(flat_p)
            return jnp.ravel(residual_fn(p))

        # Compute residuals and one reusable VJP closure at the current point.
        r, pullback = jax.vjp(flat_residual_fn, flat_params)
        damping = jnp.maximum(state.damping, 1e-12)
        rhs = -pullback(r)[0]

        def matvec(v):
            _, jv = jax.jvp(flat_residual_fn, (flat_params,), (v,))
            return pullback(jv)[0] + damping * v

        delta, cg_residual, cg_iters = _conjugate_gradient(
            matvec,
            rhs,
            max_iters=self.cg_iters,
            tol=self.cg_tol,
        )

        step_size = self.learning_rate * lr_scale
        new_flat_params = flat_params + step_size * delta
        new_params = unflatten(new_flat_params)

        new_r = flat_residual_fn(new_flat_params)
        new_loss = jnp.sum(new_r**2)

        new_state = ImplicitGaussNewtonState(
            count=state.count + 1,
            damping=state.damping,
            last_loss=new_loss,
            last_cg_residual=cg_residual,
            last_cg_iters=cg_iters,
        )
        return new_params, new_state

LevenbergMarquardt

Levenberg-Marquardt optimizer (adaptive damped Gauss-Newton).

Source code in src/deqn_jax/optimizers/gauss_newton.py 
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
class LevenbergMarquardt:
    """Levenberg-Marquardt optimizer (adaptive damped Gauss-Newton)."""

    def __init__(
        self,
        learning_rate: float = 1.0,
        initial_damping: float = 1e-3,
        damping_increase: float = 10.0,
        damping_decrease: float = 0.1,
        min_damping: float = 1e-8,
        max_damping: float = 1e8,
    ):
        self.learning_rate = learning_rate
        self.initial_damping = initial_damping
        self.damping_increase = damping_increase
        self.damping_decrease = damping_decrease
        self.min_damping = min_damping
        self.max_damping = max_damping

    def init(self, params) -> GaussNewtonState:
        return GaussNewtonState(
            count=jnp.asarray(0, dtype=jnp.int32),
            damping=self.initial_damping,
            last_loss=jnp.asarray(jnp.inf),
        )

    def update(
        self,
        residual_fn: Callable,
        params: Any,
        state: GaussNewtonState,
        lr_scale: Any = 1.0,
    ) -> Tuple[Any, GaussNewtonState]:
        """Perform one LM step with adaptive damping."""
        # Flatten params
        flat_params, unflatten = jax.flatten_util.ravel_pytree(params)
        n_params = flat_params.shape[0]

        def flat_residual_fn(flat_p):
            p = unflatten(flat_p)
            return jnp.ravel(residual_fn(p))

        # Compute residuals and Jacobian
        r = flat_residual_fn(flat_params)
        n_residuals = r.shape[0]
        current_loss = jnp.sum(r**2)

        if n_residuals <= n_params:
            J = jax.jacrev(flat_residual_fn)(flat_params)
        else:
            J = jax.jacfwd(flat_residual_fn)(flat_params)

        # Solve with current damping (dual formulation when underdetermined)
        damping = jnp.maximum(state.damping, 1e-6)
        if n_residuals < n_params:
            G = J @ J.T + damping * jnp.eye(n_residuals)
            v = jnp.linalg.solve(G, r)
            delta = -J.T @ v
        else:
            JtJ = J.T @ J + damping * jnp.eye(n_params)
            delta = jnp.linalg.solve(JtJ, -(J.T @ r))

        # Match the train_step contract used by the other optimizers: when a
        # schedule is active, self.learning_rate is 1.0 and lr_scale carries
        # the per-step learning rate.
        step_size = self.learning_rate * lr_scale
        new_flat_params = flat_params + step_size * delta
        new_r = flat_residual_fn(new_flat_params)
        new_loss = jnp.sum(new_r**2)

        # Gain ratio for damping adaptation
        Jdelta = J @ delta
        predicted = -2 * Jdelta.T @ r - Jdelta.T @ Jdelta
        actual = current_loss - new_loss
        rho = jnp.where(jnp.abs(predicted) > 1e-10, actual / predicted, 1.0)

        # Adapt damping
        new_damping = jnp.where(
            rho > 0.75,
            jnp.maximum(self.min_damping, state.damping * self.damping_decrease),
            jnp.where(
                rho < 0.25,
                jnp.minimum(self.max_damping, state.damping * self.damping_increase),
                state.damping,
            ),
        )

        # LM should only accept steps that improve the actual objective.
        accept = actual > 0.0
        final_params = jnp.where(accept, new_flat_params, flat_params)
        final_loss = jnp.where(accept, new_loss, current_loss)
        reject_damping = jnp.minimum(
            self.max_damping, state.damping * self.damping_increase
        )
        final_damping = jnp.where(accept, new_damping, reject_damping)

        new_state = GaussNewtonState(
            count=state.count + 1,
            damping=final_damping,
            last_loss=final_loss,
        )

        return unflatten(final_params), new_state

update 

update(
    residual_fn: Callable,
    params: Any,
    state: GaussNewtonState,
    lr_scale: Any = 1.0,
) -> Tuple[Any, GaussNewtonState]

Perform one LM step with adaptive damping.

Source code in src/deqn_jax/optimizers/gauss_newton.py 
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
def update(
    self,
    residual_fn: Callable,
    params: Any,
    state: GaussNewtonState,
    lr_scale: Any = 1.0,
) -> Tuple[Any, GaussNewtonState]:
    """Perform one LM step with adaptive damping."""
    # Flatten params
    flat_params, unflatten = jax.flatten_util.ravel_pytree(params)
    n_params = flat_params.shape[0]

    def flat_residual_fn(flat_p):
        p = unflatten(flat_p)
        return jnp.ravel(residual_fn(p))

    # Compute residuals and Jacobian
    r = flat_residual_fn(flat_params)
    n_residuals = r.shape[0]
    current_loss = jnp.sum(r**2)

    if n_residuals <= n_params:
        J = jax.jacrev(flat_residual_fn)(flat_params)
    else:
        J = jax.jacfwd(flat_residual_fn)(flat_params)

    # Solve with current damping (dual formulation when underdetermined)
    damping = jnp.maximum(state.damping, 1e-6)
    if n_residuals < n_params:
        G = J @ J.T + damping * jnp.eye(n_residuals)
        v = jnp.linalg.solve(G, r)
        delta = -J.T @ v
    else:
        JtJ = J.T @ J + damping * jnp.eye(n_params)
        delta = jnp.linalg.solve(JtJ, -(J.T @ r))

    # Match the train_step contract used by the other optimizers: when a
    # schedule is active, self.learning_rate is 1.0 and lr_scale carries
    # the per-step learning rate.
    step_size = self.learning_rate * lr_scale
    new_flat_params = flat_params + step_size * delta
    new_r = flat_residual_fn(new_flat_params)
    new_loss = jnp.sum(new_r**2)

    # Gain ratio for damping adaptation
    Jdelta = J @ delta
    predicted = -2 * Jdelta.T @ r - Jdelta.T @ Jdelta
    actual = current_loss - new_loss
    rho = jnp.where(jnp.abs(predicted) > 1e-10, actual / predicted, 1.0)

    # Adapt damping
    new_damping = jnp.where(
        rho > 0.75,
        jnp.maximum(self.min_damping, state.damping * self.damping_decrease),
        jnp.where(
            rho < 0.25,
            jnp.minimum(self.max_damping, state.damping * self.damping_increase),
            state.damping,
        ),
    )

    # LM should only accept steps that improve the actual objective.
    accept = actual > 0.0
    final_params = jnp.where(accept, new_flat_params, flat_params)
    final_loss = jnp.where(accept, new_loss, current_loss)
    reject_damping = jnp.minimum(
        self.max_damping, state.damping * self.damping_increase
    )
    final_damping = jnp.where(accept, new_damping, reject_damping)

    new_state = GaussNewtonState(
        count=state.count + 1,
        damping=final_damping,
        last_loss=final_loss,
    )

    return unflatten(final_params), new_state

gauss_newton 

gauss_newton(
    learning_rate: float = 1.0, damping: float = 0.0
) -> GaussNewton

Create Gauss-Newton optimizer.

Parameters:

Name Type Description Default
learning_rate float

Step size multiplier (1.0 = full GN step)

 1.0
damping float

Fixed damping (0 = pure GN, >0 = LM-style regularization)

 0.0

Returns:

Type Description
GaussNewton

GaussNewton optimizer instance
Source code in src/deqn_jax/optimizers/gauss_newton.py 
288
289
290
291
292
293
294
295
296
297
298
299
300
301
def gauss_newton(
    learning_rate: float = 1.0,
    damping: float = 0.0,
) -> GaussNewton:
    """Create Gauss-Newton optimizer.

    Args:
        learning_rate: Step size multiplier (1.0 = full GN step)
        damping: Fixed damping (0 = pure GN, >0 = LM-style regularization)

    Returns:
        GaussNewton optimizer instance
    """
    return GaussNewton(learning_rate, damping)

implicit_gauss_newton 

implicit_gauss_newton(
    learning_rate: float = 1.0,
    damping: float = 0.0001,
    cg_iters: int = 20,
    cg_tol: float = 1e-06,
) -> ImplicitGaussNewton

Create a matrix-free damped Gauss-Newton optimizer.

Source code in src/deqn_jax/optimizers/gauss_newton.py 
304
305
306
307
308
309
310
311
312
313
314
315
316
317
def implicit_gauss_newton(
    learning_rate: float = 1.0,
    damping: float = 1e-4,
    cg_iters: int = 20,
    cg_tol: float = 1e-6,
) -> ImplicitGaussNewton:
    """Create a matrix-free damped Gauss-Newton optimizer."""

    return ImplicitGaussNewton(
        learning_rate=learning_rate,
        damping=damping,
        cg_iters=cg_iters,
        cg_tol=cg_tol,
    )

levenberg_marquardt 

levenberg_marquardt(
    learning_rate: float = 1.0,
    initial_damping: float = 0.001,
    damping_increase: float = 10.0,
    damping_decrease: float = 0.1,
    min_damping: float = 1e-08,
    max_damping: float = 100000000.0,
) -> LevenbergMarquardt

Create Levenberg-Marquardt optimizer (adaptive damped GN).

Parameters:

Name Type Description Default
learning_rate float

Step size multiplier

 1.0
initial_damping float

Starting damping value

 0.001
damping_increase float

Factor when step is bad

 10.0
damping_decrease float

Factor when step is good

 0.1
min_damping float

Minimum damping

 1e-08
max_damping float

Maximum damping

 100000000.0

Returns:

Type Description
LevenbergMarquardt

LevenbergMarquardt optimizer instance
Source code in src/deqn_jax/optimizers/gauss_newton.py 
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
def levenberg_marquardt(
    learning_rate: float = 1.0,
    initial_damping: float = 1e-3,
    damping_increase: float = 10.0,
    damping_decrease: float = 0.1,
    min_damping: float = 1e-8,
    max_damping: float = 1e8,
) -> LevenbergMarquardt:
    """Create Levenberg-Marquardt optimizer (adaptive damped GN).

    Args:
        learning_rate: Step size multiplier
        initial_damping: Starting damping value
        damping_increase: Factor when step is bad
        damping_decrease: Factor when step is good
        min_damping: Minimum damping
        max_damping: Maximum damping

    Returns:
        LevenbergMarquardt optimizer instance
    """
    return LevenbergMarquardt(
        learning_rate,
        initial_damping,
        damping_increase,
        damping_decrease,
        min_damping,
        max_damping,
    )

make_grad_step_gn 

make_grad_step_gn(
    model,
    opt: Any,
    mc_samples: int,
    batch_size: int,
    quad_nodes,
    quad_weights,
    loss_reweight: str,
    reweight_alpha: float,
    use_target_network: bool,
    compute_loss_fn,
)

JIT'd: one Gauss-Newton / Levenberg-Marquardt update on a minibatch.

Source code in src/deqn_jax/optimizers/gauss_newton.py 
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
def make_grad_step_gn(
    model,
    opt: Any,
    mc_samples: int,
    batch_size: int,
    quad_nodes,
    quad_weights,
    loss_reweight: str,
    reweight_alpha: float,
    use_target_network: bool,
    compute_loss_fn,
):
    """JIT'd: one Gauss-Newton / Levenberg-Marquardt update on a minibatch."""
    import equinox as eqx
    import optax

    from deqn_jax.training.loss import (
        compute_loss,
        compute_residuals,
        sample_antithetic_shocks,
    )
    from deqn_jax.training.reweighting import update_reweighting
    from deqn_jax.types import Metrics

    n_eq = len(model.equation_names) if model.equation_names else 1
    _compute_loss_log = compute_loss_fn or compute_loss
    use_quadrature = quad_nodes is not None and quad_weights is not None

    @jax.jit
    def grad_step(
        state,
        batch,
        lr_scale,
        shock_scale=jnp.array(1.0),
    ):
        loss_key, new_key = jax.random.split(state.key)
        target_fn = state.target_params if use_target_network else None

        def residual_fn(params):
            if use_quadrature:
                n_nodes = quad_nodes.shape[0]
                shocks = (
                    jnp.broadcast_to(
                        quad_nodes[:, None, :],
                        (n_nodes, batch_size, model.n_shocks),
                    )
                    * shock_scale
                )
                sample_weights = quad_weights
            else:
                shocks = sample_antithetic_shocks(
                    loss_key,
                    mc_samples,
                    batch_size,
                    model.n_shocks,
                    shock_scale,
                )
                n_samples = shocks.shape[0]
                sample_weights = jnp.ones(n_samples) / n_samples

            def sample_residuals(shock):
                return compute_residuals(
                    model, params, batch, shock, target_policy_fn=target_fn
                )

            all_residuals = jax.vmap(sample_residuals)(shocks)
            per_eq = []
            for r in all_residuals.values():
                mean_r = jnp.einsum("s,sb->b", sample_weights, r)
                per_eq.append(mean_r)
            return jnp.concatenate(per_eq)

        loss, eq_losses = _compute_loss_log(
            model,
            state.params,
            batch,
            loss_key,
            mc_samples,
            weights=state.loss_weights,
            shock_scale=shock_scale,
            quad_nodes=quad_nodes,
            quad_weights=quad_weights,
            target_policy_fn=target_fn,
        )

        new_params, new_opt_state = opt.update(
            residual_fn, state.params, state.opt_state, lr_scale=lr_scale
        )

        def scalar_loss(p):
            r = residual_fn(p)
            return jnp.sum(r**2)

        grad_norm = optax.global_norm(
            eqx.filter(jax.grad(scalar_loss)(state.params), eqx.is_array)
        )

        new_weights, new_rw = update_reweighting(
            eq_losses,
            state,
            loss_reweight,
            reweight_alpha,
            n_eq,
        )
        new_state = state._replace(
            params=new_params,
            opt_state=new_opt_state,
            key=new_key,
            step=state.step + 1,
            loss_weights=new_weights,
            reweight_state=new_rw,
        )
        return new_state, Metrics(loss=loss, residuals=eq_losses, grad_norm=grad_norm)

    return grad_step