Types

The framework's core types are NamedTuples for JAX pytree compatibility. ModelSpec is the user contract — the only object a model needs to populate to be trainable. TrainState carries everything mutable across a training run (params, opt state, episode state, key, step counter, replay state, …) so train_step can be a pure function suitable for @jax.jit.

For the ModelSpec field listing with signatures and shape contracts, see The user contract in REFERENCE.md. For the prose walkthrough on populating a ModelSpec, see Implementing a model.

The remaining types (ReweightState, ReplayState, EpisodeState, Metrics) are framework-internal data plumbing that you'll see in type signatures but rarely construct yourself. The exception is make_reweight_state(n_equations) — used when constructing a TrainState outside create_train_state (rare; mainly in tests).

deqn_jax.types

Core type definitions for DEQN-JAX.

Uses NamedTuples for pytree compatibility with JAX transformations.

ModelSpec

Bases: NamedTuple

Specification for an economic model.

A ModelSpec defines everything needed to train a DEQN: - Dimensions (states, policies, shocks) - Economic equations (equilibrium conditions) - State transitions (dynamics) - Constants (calibration parameters)

All functions should be pure and JAX-compatible.

Example

MODEL = ModelSpec( name="brock_mirman", n_states=2, n_policies=1, n_shocks=1, state_names=["k", "z"], policy_names=["sav_rate"], constants={"alpha": 1/3, "beta": 0.95, ...}, equations_fn=equations, step_fn=step, steady_state_fn=compute_steady_state, )

Source code in src/deqn_jax/types.py

class ModelSpec(NamedTuple):
    """Specification for an economic model.

    A ModelSpec defines everything needed to train a DEQN:
    - Dimensions (states, policies, shocks)
    - Economic equations (equilibrium conditions)
    - State transitions (dynamics)
    - Constants (calibration parameters)

    All functions should be pure and JAX-compatible.

    Example:
        MODEL = ModelSpec(
            name="brock_mirman",
            n_states=2,
            n_policies=1,
            n_shocks=1,
            state_names=["k", "z"],
            policy_names=["sav_rate"],
            constants={"alpha": 1/3, "beta": 0.95, ...},
            equations_fn=equations,
            step_fn=step,
            steady_state_fn=compute_steady_state,
        )
    """

    name: str
    n_states: int
    n_policies: int
    n_shocks: int

    # Function signatures:
    # equations_fn(state, policy, next_state, next_policy, constants) -> Dict[str, Array]
    equations_fn: Callable[..., Dict[str, Array]]

    # step_fn(state, policy, shock, constants) -> next_state
    step_fn: Callable[..., Array]

    # Model constants (calibration parameters)
    constants: Dict[str, float]

    # Variable name metadata. Defaults to empty tuple so the framework
    # can iterate / index / list() unconditionally without None-checks.
    # Every shipped model populates them; toy ModelSpecs in tests fall
    # back to () with no behaviour change (truthy checks treat () like
    # the legacy None and ``list(()) == []``).
    state_names: Tuple[str, ...] = ()
    policy_names: Tuple[str, ...] = ()
    equation_names: Tuple[str, ...] = ()

    # Optional: compute steady state for warm-starting
    # steady_state_fn(constants) -> (ss_state, ss_policy)
    steady_state_fn: Optional[Callable[..., Tuple[Array, Array]]] = None

    # Optional: initial state distribution sampler
    # init_state_fn(key, batch_size, constants) -> state
    init_state_fn: Optional[Callable[..., Array]] = None

    # Optional: compute derived quantities for monitoring
    # definitions_fn(state, policy, constants) -> Dict[str, Array]
    definitions_fn: Optional[Callable[..., Dict[str, Any]]] = None

    # Optional two-stage residual hooks for models whose residual wraps an
    # EXPECTATION in a nonlinearity -- e.g. a Fischer-Burmeister borrowing
    # constraint on an intertemporal Euler, where E[.] must be taken BEFORE the
    # FB. When ``combine_fn`` is set, the loss:
    #   1. evaluates ``inside_fn`` per shock and averages it -> E[inside]
    #      (MC-safe: ``inside`` must be linear in the shock-dependent terms),
    #   2. applies ``combine_fn`` to produce the final per-equation residuals
    #      (the nonlinearity, applied AFTER the expectation).
    # Standard models leave both None and use ``equations_fn``; the existing
    # ``(E[residual])^2`` path is exactly the special case combine = identity.
    #   inside_fn(state, policy, next_state, next_policy, constants) -> Dict[str, Array]
    #   combine_fn(state, policy, E_inside: Dict[str, Array], constants) -> Dict[str, Array]
    inside_fn: Optional[Callable[..., Dict[str, Array]]] = None
    combine_fn: Optional[Callable[..., Dict[str, Array]]] = None

    # Policy bounds (for output activation)
    policy_lower: Optional[Array] = None
    policy_upper: Optional[Array] = None

    # Per-policy default output parameterization for residual networks (e.g.
    # ``LinearPlusMLP``, ``DisasterPolicyNet``). Each entry "linear" or "log".
    # ``"log"`` means the network outputs π_i = ss_i * exp(BK_log + MLP),
    # baking in positivity. Length must equal len(policy_names). When None,
    # networks default to all-linear (legacy behavior). User-supplied
    # ``network.output_links`` in YAML overrides this.
    default_output_links: Optional[Tuple[str, ...]] = None

    # Discrete Markov-chain shock support. When set, the trainer:
    #   - rollout: samples next z-index from Multinomial(transition_matrix[z_t]),
    #   - residual expectation: enumerates over all K next-z values, weighting
    #     residuals by transition_matrix[z_t, z_{t+1}].
    # Activated by ``TrainConfig.expectation_type='discrete'``. Models that
    # encode a finite-cardinality exogenous chain (e.g. discrete TFP, OLG
    # depreciation states) should provide both fields below; the shock passed
    # to ``step_fn`` is then a categorical index in [0, K), not a continuous
    # noise. Models without these fields use Gaussian shocks as before.
    transition_matrix: Optional[Array] = None  # [K, K], rows sum to 1
    z_state_idx: Optional[int] = None  # which entry of state holds current z-index

    # Optional: clip states for simulation safety (eval/irf only, NOT training)
    # clip_state_fn(state) -> state
    clip_state_fn: Optional[Callable[..., Array]] = None

    # Optional: box barrier penalty on states (added to loss)
    # state_barrier_fn(state) -> penalty [batch]
    state_barrier_fn: Optional[Callable[..., Array]] = None

    # Optional: shock names for diagnostics/logging
    shock_names: Optional[Tuple[str, ...]] = None

    # Optional: called every ``log_every`` episodes during training, after
    # scalar/histogram logging. Side-effect only (writes plots, logs to TB,
    # etc.). Signature:
    #     cycle_hook(state: TrainState, model: ModelSpec, episode: int) -> None
    # The hook should close over any configuration it needs (output dir,
    # logger, etc.) at model-construction time. See DEQN-MAO's
    # model-level ``Hooks.py`` for the reference pattern; our version
    # differs only in that the plotting primitives themselves live in the
    # shared ``deqn_jax.plots`` module and the hook composes them.
    cycle_hook: Optional[Callable[..., None]] = None

    # Optional declarative bound specs. Format (matches DEQN-MAO upstream):
    #     {"name": {"lower": float, "upper": float,
    #               "penalty_lower": float, "penalty_upper": float}}
    # When set, the loss picks up a soft-penalty term
    #     penalty_lower * mean(max(0, lower - value) ** 2)
    # for each bounded variable (analogous for upper). Missing penalty
    # coefficients default to 1/bound**2 (upstream convention). Use
    # state_bounds for states and definition_bounds for derived quantities
    # computed via ``definitions_fn``. Hard-clipped policies are enforced
    # via ``policy_lower``/``policy_upper`` with the activation layer,
    # separately from this soft mechanism.
    state_bounds: Optional[Dict[str, Dict[str, float]]] = None
    definition_bounds: Optional[Dict[str, Dict[str, float]]] = None

    # Optional: called once before training starts, given the loaded
    # model and the resolved TrainConfig. Returns a (possibly modified)
    # ModelSpec for the trainer to use. Use this to wire config-time
    # decisions into the model -- e.g. disaster swaps `steady_state_fn`
    # to its risky-SS variant when ``constants['p_disaster'] > 0`` and
    # ``config.use_risky_steady_state`` allows it. Called outside JIT,
    # so plain Python branching is fine. Default ``None`` is a no-op
    # (model is used as-declared).
    #
    # Signature: setup_fn(model: ModelSpec, config) -> ModelSpec
    setup_fn: Optional[Callable[..., "ModelSpec"]] = None

    # Optional: called every ``log_every`` cycles in the Python-level
    # logging path, given the model and current training-batch quantities,
    # to return a dict of scalar diagnostics that the trainer prepends
    # to TensorBoard / W&B with the model's namespace prefix. Lets a
    # model expose its own per-equation decompositions, ratio
    # diagnostics, soft-floor saturation fractions, etc. without the
    # framework knowing about the model's internals. Failure is
    # tolerated -- if the hook raises, the trainer logs a warning and
    # continues.
    #
    # Signature:
    #   scalar_diagnostics_fn(
    #       model: ModelSpec,
    #       policy_fn: Callable,        # eqx Module or wrapper for sequence nets
    #       states: Array,              # [batch, n_states] training batch
    #       policy_out: Array,          # [batch, n_policies] policy at states
    #       defs: Dict[str, Array],     # definitions at (states, policy_out)
    #   ) -> Dict[str, float]
    scalar_diagnostics_fn: Optional[Callable[..., Dict[str, float]]] = None

    # Optional: model-specific auxiliary terms for the composite loss
    # (``loss_type='composite'``). Lets a model contribute extra
    # ``aux_*``-keyed losses without the framework knowing about that
    # model's definitions or solver internals. Called inside
    # ``make_composite_loss``'s closure after barrier losses, with the
    # batch-level ``defs`` dict and a kwargs dict of relevant
    # ``CompositeLossConfig`` weights. Returns ``(aux_entries,
    # total_contribution)``: ``aux_entries`` is merged into
    # ``eq_losses`` (so adaptive reweighting / logging see the
    # individual unweighted scalars under their ``aux_*`` keys);
    # ``total_contribution`` is added directly to the running total
    # (the hook applies its own weighting).
    #
    # Used by the disaster model to add Newton-step diagnostic losses
    # (``aux_newton_cond``, ``aux_newton_resid``) that read disaster-
    # specific definitions (``newton_h_prime``, ``newton_residual``).
    # Other models leave this ``None``.
    #
    # Signature:
    #   composite_aux_fn(
    #       model: ModelSpec,
    #       defs: Dict[str, Array],     # batch-level definitions
    #       data,                       # CompositeData (linearization + SS)
    #       weights: Dict[str, float],  # subset of CompositeLossConfig weights
    #   ) -> Tuple[Dict[str, Array], Array]
    composite_aux_fn: Optional[Callable[..., Tuple[Dict[str, Array], Array]]] = None

    # Optional: precompute model-specific constants for the composite-loss
    # aux hook (e.g. SS leverage for the disaster barrier threshold). Called
    # once at trainer setup, OUTSIDE jit; the returned dict lands in
    # ``CompositeData.aux_constants`` and is read by ``composite_aux_fn`` at
    # loss-evaluation time. Default ``None`` means no aux constants needed.
    #
    # Signature:
    #   composite_aux_constants_fn(model: ModelSpec) -> Dict[str, Any]
    composite_aux_constants_fn: Optional[Callable[..., Dict[str, Any]]] = None

ReweightState

Bases: NamedTuple

Running statistics for adaptive loss reweighting.

Used by lr_annealing and relobralo strategies to track per-equation loss history for dynamic weight adjustment.

Attributes:

Name	Type	Description
running_ema	Array	EMA of per-equation losses (for lr_annealing)
prev_losses	Array	Previous step losses (for relobralo)
init_losses	Array	First step losses (for relobralo)
initialized	Array	Whether init_losses has been set

Source code in src/deqn_jax/types.py

class ReweightState(NamedTuple):
    """Running statistics for adaptive loss reweighting.

    Used by lr_annealing and relobralo strategies to track
    per-equation loss history for dynamic weight adjustment.

    Attributes:
        running_ema: EMA of per-equation losses (for lr_annealing)
        prev_losses: Previous step losses (for relobralo)
        init_losses: First step losses (for relobralo)
        initialized: Whether init_losses has been set
    """

    running_ema: Array  # [n_eq]
    prev_losses: Array  # [n_eq]
    init_losses: Array  # [n_eq]
    initialized: Array  # scalar bool

ReplayState

Bases: NamedTuple

Fixed-shape ring buffer of past states + per-state priorities.

Used by the prioritized state-replay mechanism (RL-style anti-forgetting + spectral-bias mitigation). Lives on TrainState so it survives checkpoint resume and is variant-agnostic across the 5 train-step kinds.

Attributes:

Name	Type	Description
buffer	Array	Past states [capacity, n_states], fp32. Newest writes land at write_idx; the slot is overwritten on wrap-around.
priorities	Array	Per-row scalar priorities ([capacity], fp32). Computed at write time as the per-element sum of squared equilibrium residuals; sampling probability is proportional to (priority + eps) ** alpha.
write_idx	Array	Scalar int32 ring cursor. Modulo capacity after writes.
n_filled	Array	Scalar int32, capped at capacity. Used as the slice upper bound for sampling so unfilled tail rows are never drawn.

Source code in src/deqn_jax/types.py

class ReplayState(NamedTuple):
    """Fixed-shape ring buffer of past states + per-state priorities.

    Used by the prioritized state-replay mechanism (RL-style anti-forgetting +
    spectral-bias mitigation). Lives on ``TrainState`` so it survives
    checkpoint resume and is variant-agnostic across the 5 train-step kinds.

    Attributes:
        buffer: Past states ``[capacity, n_states]``, fp32. Newest writes
            land at ``write_idx``; the slot is overwritten on wrap-around.
        priorities: Per-row scalar priorities (``[capacity]``, fp32).
            Computed at write time as the per-element sum of squared
            equilibrium residuals; sampling probability is proportional
            to ``(priority + eps) ** alpha``.
        write_idx: Scalar int32 ring cursor. Modulo ``capacity`` after writes.
        n_filled: Scalar int32, capped at ``capacity``. Used as the slice
            upper bound for sampling so unfilled tail rows are never drawn.
    """

    buffer: Array  # [capacity, n_states]
    priorities: Array  # [capacity]
    write_idx: Array  # scalar int32
    n_filled: Array  # scalar int32

TrainState

Bases: NamedTuple

Immutable training state for JAX-compatible training loops.

All state needed for training is bundled here so that train_step can be a pure function suitable for jax.jit.

Attributes:

Name	Type	Description
params	Any	Equinox model (pytree of parameters)
opt_state	Any	Optax optimizer state
episode_state	Array	Current state batch for episode simulation
key	Array	JAX PRNG key
step	int	Current training step
episode	int	Current episode number
loss_weights	Array	Per-equation loss weights [n_eq]
reweight_state	ReweightState	Adaptive reweighting running statistics
target_params	Any	Frozen policy copy for target-network style training
aux_params	Any	Slot for a second trainable module (e.g. value network in actor-critic, critic network, learned expectation operator). None by default. Default training loop ignores it; only loss functions that know about aux_params will use it.
aux_opt_state	Any	Optimizer state for aux_params if trained with its own optimizer. None if aux is trained jointly with the primary optimizer.
history_state	Any	Sliding history window [batch, H, n_states] for sequence policies (LSTM/Transformer, network.history_len > 1). Persists across rollouts so recurrent training sees continuous ergodic trajectories rather than rebuilding a constant window at every cycle. None for MLP models (history_len == 1).
replay_state	Any	Prioritized state-replay buffer (ReplayState). Off by default (None). When enabled via TrainConfig.replay_buffer.enabled, holds a fixed-shape ring buffer of past trajectory states + their write-time residual priorities; the cycle hooks read/write it once per cycle to mix buffered samples into each minibatch dataset.

Source code in src/deqn_jax/types.py

class TrainState(NamedTuple):
    """Immutable training state for JAX-compatible training loops.

    All state needed for training is bundled here so that train_step
    can be a pure function suitable for jax.jit.

    Attributes:
        params: Equinox model (pytree of parameters)
        opt_state: Optax optimizer state
        episode_state: Current state batch for episode simulation
        key: JAX PRNG key
        step: Current training step
        episode: Current episode number
        loss_weights: Per-equation loss weights [n_eq]
        reweight_state: Adaptive reweighting running statistics
        target_params: Frozen policy copy for target-network style training
        aux_params: Slot for a second trainable module (e.g. value network
            in actor-critic, critic network, learned expectation operator).
            None by default. Default training loop ignores it; only loss
            functions that know about ``aux_params`` will use it.
        aux_opt_state: Optimizer state for ``aux_params`` if trained with
            its own optimizer. ``None`` if aux is trained jointly with the
            primary optimizer.
        history_state: Sliding history window ``[batch, H, n_states]`` for
            sequence policies (LSTM/Transformer, ``network.history_len > 1``).
            Persists across rollouts so recurrent training sees continuous
            ergodic trajectories rather than rebuilding a constant window
            at every cycle. ``None`` for MLP models (``history_len == 1``).
        replay_state: Prioritized state-replay buffer (``ReplayState``).
            Off by default (``None``). When enabled via
            ``TrainConfig.replay_buffer.enabled``, holds a fixed-shape
            ring buffer of past trajectory states + their write-time
            residual priorities; the cycle hooks read/write it once per
            cycle to mix buffered samples into each minibatch dataset.
    """

    params: Any  # Equinox model
    opt_state: Any  # Optax optimizer state
    episode_state: Array  # Current states [batch, n_states]
    key: Array  # JAX PRNG key
    step: int
    episode: int
    loss_weights: Array  # [n_eq] per-equation weights
    reweight_state: ReweightState  # adaptive reweighting state
    target_params: Any = None  # Frozen policy for target network (DQN-style)
    aux_params: Any = None  # Auxiliary trainable module (value net, critic, ...)
    aux_opt_state: Any = None  # Optimizer state for aux_params if separate
    history_state: Any = None  # [batch, H, n_states] for sequence policies, else None
    replay_state: Any = None  # ReplayState (prioritized state buffer); None when off

EpisodeState

Bases: NamedTuple

State carried through episode simulation via lax.scan.

Attributes:

Name	Type	Description
state	Array	Current economic state [batch, n_states]
key	Array	PRNG key for shock sampling

Source code in src/deqn_jax/types.py

class EpisodeState(NamedTuple):
    """State carried through episode simulation via lax.scan.

    Attributes:
        state: Current economic state [batch, n_states]
        key: PRNG key for shock sampling
    """

    state: Array
    key: Array

Metrics

Bases: NamedTuple

Training metrics from a single step/episode.

All three fields hold scalar JAX Arrays at runtime (built inside JIT'd grad steps). They were previously annotated as plain Python float / Dict[str, float], which produced ~70 spurious reportArgumentType errors at every Metrics(...) call site -- consumers cast to float explicitly when they need a Python scalar (float(metrics.loss)).

Source code in src/deqn_jax/types.py

class Metrics(NamedTuple):
    """Training metrics from a single step/episode.

    All three fields hold scalar JAX Arrays at runtime (built inside
    JIT'd grad steps). They were previously annotated as plain Python
    ``float`` / ``Dict[str, float]``, which produced ~70 spurious
    ``reportArgumentType`` errors at every Metrics(...) call site --
    consumers cast to ``float`` explicitly when they need a Python
    scalar (``float(metrics.loss)``).
    """

    loss: Array
    residuals: Optional[Dict[str, Array]] = None
    grad_norm: Optional[Array] = None

make_reweight_state

make_reweight_state(n_equations: int) -> ReweightState

Create initial reweight state for n equations.

Source code in src/deqn_jax/types.py

def make_reweight_state(n_equations: int) -> "ReweightState":
    """Create initial reweight state for n equations."""
    return ReweightState(
        running_ema=jnp.ones(n_equations),
        prev_losses=jnp.zeros(n_equations),
        init_losses=jnp.zeros(n_equations),
        initialized=jnp.array(False),
    )

make_replay_state

make_replay_state(
    capacity: int, n_states: int
) -> ReplayState

Create an empty replay buffer with room for capacity states.

Source code in src/deqn_jax/types.py

def make_replay_state(capacity: int, n_states: int) -> "ReplayState":
    """Create an empty replay buffer with room for ``capacity`` states."""
    return ReplayState(
        buffer=jnp.zeros((capacity, n_states), dtype=jnp.float32),
        priorities=jnp.zeros((capacity,), dtype=jnp.float32),
        write_idx=jnp.array(0, dtype=jnp.int32),
        n_filled=jnp.array(0, dtype=jnp.int32),
    )