Skip to content

Quickstart

Train the canonical model with the validated stack, then read its accuracy the way you'd report it in a paper — the relative-Euler-error (errREE) distribution on the ergodic path.

Status: alpha (v0.2.0) — the validated stack is small

Everything below leads with the combination the test suite and gallery actually exercise: adam + an MLP + an MSE residual + antithetic Monte-Carlo expectations, on brock_mirman. Everything else in the registries is a research instrument, not a turnkey recommendation — see the Method Zoo for when (and when not) to reach for it.

0. Verify the install

uv sync
uv run deqn-jax check     # JAX backend, devices, registered models & optimizers
uv run deqn-jax list      # the registered models
Install detail — source checkout, CUDA, editable mode

Alpha is not yet on PyPI; install from a source checkout.

git clone <repo>
cd deqn-jax
uv sync
uv pip install -e .            # optional: editable mode for hacking

GPU build (Linux aarch64 / x86_64):

uv pip install -U "jax[cuda13]"   # or "jax[cuda12]" for CUDA 12

uv run deqn-jax check reports the active backend and devices. Always use uv run; never activate the venv by hand.

1. Solve a model in five minutes

brock_mirman is the canonical/teaching tier: state \((k, z)\), one decision rule (the savings rate), one consumption Euler equation, an analytical steady state. It is the smoke test that proves the stack works on your machine.

uv run deqn-jax train brock_mirman -n 1000 --warm-start \
    --checkpoint-dir checkpoints/brock_mirman

You should see the residual loss fall several orders of magnitude in under a minute on CPU. --warm-start fits the network to the steady-state policy first (an L-BFGS supervised pre-fit), so training begins from a sane economic guess rather than noise; --checkpoint-dir is what lets the next step read the trained policy back.

What just happened, in your language

The network plays the role Chebyshev polynomials or splines play in a projection method — a flexible approximation of the decision rule \(\pi(s)\). "Training" is the collocation/projection solve for its coefficients; "minibatches" are collocation states drawn by simulating the model (the ergodic set), not a fixed tensor grid; the "loss" is the Euler residual, integrated over next-period shocks by antithetic Monte Carlo.

2. Read its accuracy

A low loss is necessary but not sufficient — like any nonlinear global solver, residual-minimization can land on the wrong answer. So you don't trust the loss; you check the policy. evaluate simulates a long ergodic path and reports the errREE distribution — the gold-standard accuracy metric (Azinovic et al. 2022), the number you'd quote.

uv run deqn-jax evaluate checkpoints/brock_mirman/checkpoint_best.eqx -n 10000

It also runs the market-clearing, simulated-moments, and stability checks; the config is auto-detected from the checkpoint directory. For measured errREE certificates on worked models, see the Gallery — the evidence, not a promise.

3. The everyday loop

uv run deqn-jax train brock_mirman -n 1000 --warm-start \
    --checkpoint-dir checkpoints/brock_mirman

Swap in any registered model: bm_labor_constrained, irbc, olg_lifecycle. Config-driven runs read a YAML and accept dot-notation overrides:

uv run deqn-jax train --config configs/brock_mirman.yaml \
    --set optimizer.learning_rate=0.001 \
    --checkpoint-dir checkpoints/brock_mirman
uv run deqn-jax evaluate checkpoints/brock_mirman/checkpoint_best.eqx -n 10000

The errREE distribution, market-clearing errors, simulated moments, and the stability gate. The config is auto-detected from the checkpoint directory.

uv run deqn-jax irf checkpoints/brock_mirman/checkpoint_best.eqx --shock eps_z

Impulse responses from a trained policy. --girf gives the generalized (state-dependent, no-shock-baseline-subtracted) variant for nonlinear models. Run deqn-jax info brock_mirman for valid shock names.

Where to next

  • See the sell, measured


    Closed-form pedagogy → the occasionally-binding constraint trilogy (bm_labor_constrained, irbc, olg_lifecycle) → an experimental NK-DSGE, each with its measured errREE certificate.

    Gallery

  • Pick your method


    The swappable toolkit — networks, optimizers, expectations, diagnostics — and when to reach for each. The default recipe is on the first screen.

    Method Zoo

  • Write your own model


    Declare states, equilibrium residuals, transition, calibration — as data. The ModelSpec contract is the whole surface.

    Implementing a model

Resume a checkpoint — and the Adam → Newton-style polish

Any checkpoint resumes, including with a different optimizer. The legitimate use is the pipeline the Method Zoo routes you to when a first-order run plateaus: rough exploration with adam, then a Newton-style polish — the same quasi-Newton / Gauss-Newton machinery you know from GMM / MLE estimation, applied to the equilibrium residuals for quadratic convergence near a solution.

# Rough exploration with Adam (the validated first-order method)
uv run deqn-jax train brock_mirman -n 1000 --warm-start \
    --checkpoint-dir checkpoints/brock_mirman

# Polish from the checkpoint with L-BFGS (Newton-style; experimental)
uv run deqn-jax train brock_mirman -n 200 \
    --resume checkpoints/brock_mirman/checkpoint_best.eqx \
    --set optimizer.name=lbfgs \
    --checkpoint-dir checkpoints/brock_mirman

The trainer detects the optimizer change, re-initializes optimizer state for the new method, and keeps the network weights; the original config is read from <checkpoint_dir>/config.yaml to reconstruct the pytree template. gn / lm (Gauss-Newton, Levenberg-Marquardt) are the other Newton-style members. These are experimental polish steps — adam remains the validated workhorse, and a stall is more often a network fix (linear_plus_mlp, the Blanchard-Kahn-anchored basis) than an optimizer one.

The disaster model is experimental — under validation

disaster (CMR-style NK-DSGE, 13 states / 11 policies, numerical steady state) is the stress test, not part of the validated stack. The baseline block converges, but the disaster/financial-frictions block is still under validation, and the recipe it leans on — LinearPlusMLP plus the composite loss (anchor + Jacobian-match + barrier + Newton auxiliary terms) — is itself experimental. Treat it as a research example, not a turnkey result.

uv run deqn-jax train --config configs/disaster.yaml   # experimental

See the gallery landing and the composite loss note before trusting any number it produces.

Two honest limits — stated here, not buried

  • A low residual does not pin down the right equilibrium. Like any nonlinear global solver, DEQN can settle on the wrong branch, and nothing here enforces equilibrium selection. There is no global analogue of the local Blanchard-Kahn saddle-path condition — BK is a linear/local determinacy criterion, not a global one.
  • No certified error bounds. Accuracy is measured (the errREE distribution), not proven by a theorem. Quote the number; don't assume it.