Skip to content

Models

Two registration paths share one _MODELS dict:

  • In-tree — add an import + entry to _MODELS and _DESCRIPTIONS in src/deqn_jax/models/__init__.py. Right thing for models that ship with the library.
  • Programmatic — call register_model(spec, description=...) at runtime. Right thing for agent-codegen'd models in user projects, notebook prototypes, or external plugins. See Adding a model for the contract.

Both paths feed load_model(name) and list_models() identically.

VariableSpec (in variable_spec.py) is a small helper that gives named attribute access to state and policy arrays (s.k, p.sav_rate) across both batched [batch, n] and unbatched [n] shapes — cleaner than state[:, 0] everywhere, and traces through jax.vmap unchanged. make_init_state_fn is a declarative builder for initial-state samplers; supports uniform, normal, lognormal, truncated_normal, constant per variable.

deqn_jax.models

Economic models for DEQN-JAX.

Each model subpackage exports a MODEL: ModelSpec constant. Two registration paths, both supported and orthogonal:

  1. In-tree (this file's _MODELS / _DESCRIPTIONS dicts): add an import + entry. Right thing for models that ship with the library and are version-controlled here.
  2. Programmatic (register_model(spec, description=...)): add at runtime from anywhere. Right thing for agent-codegen'd models in user projects, notebook prototypes, or pluggable extensions that don't want to fork the library.

Both paths feed the same _MODELS dict, so load_model and list_models see them identically.

load_model 

load_model(name: str) -> ModelSpec

Return the registered model named name.

Raises ValueError with the available names if not found.

Source code in src/deqn_jax/models/__init__.py 
59
60
61
62
63
64
65
66
def load_model(name: str) -> ModelSpec:
    """Return the registered model named ``name``.

    Raises ``ValueError`` with the available names if not found.
    """
    if name not in _MODELS:
        raise ValueError(f"Unknown model: {name!r}. Available: {list(_MODELS)}")
    return _MODELS[name]

list_models 

list_models() -> List[Tuple[str, str]]

Return [(name, description), ...] for every registered model.

Sees both in-tree and runtime-registered entries.

Source code in src/deqn_jax/models/__init__.py 
69
70
71
72
73
74
def list_models() -> List[Tuple[str, str]]:
    """Return ``[(name, description), ...]`` for every registered model.

    Sees both in-tree and runtime-registered entries.
    """
    return [(name, _DESCRIPTIONS.get(name, "")) for name in _MODELS]

register_model 

register_model(
    spec: ModelSpec,
    *,
    description: Optional[str] = None,
    overwrite: bool = False,
) -> None

Register a model at runtime so load_model(spec.name) finds it.

Equivalent to adding an entry to _MODELS and _DESCRIPTIONS in this file, but at import / instantiation time. Intended for:

  • agent-codegen'd models in user projects (no fork required),
  • notebook prototyping,
  • plugin packages that add models to deqn-jax via register_model on import.

Parameters:

Name Type Description Default
spec ModelSpec

A populated ModelSpec. Its spec.name field is the registry key. The framework consumes this object via the same load_model lookup as in-tree models.

 required
description Optional[str]

Human-readable one-liner for list_models. Defaults to an empty string if omitted.

 None
overwrite bool

If True, an existing entry with the same name is replaced. If False (default) and spec.name is already registered, raises ValueError so you don't shadow an in-tree model by accident.

 False

Raises:

Type Description
ValueError

If spec.name is empty, or already registered and overwrite=False.
TypeError

If spec is not a ModelSpec.

Example::

from deqn_jax.api import ModelSpec, register_model

MY_MODEL = ModelSpec(name="my_model", ...)
register_model(MY_MODEL, description="My agent-built model")

# Now usable everywhere:
from deqn_jax.api import load_model, train_from_config, TrainConfig
cfg = TrainConfig(model="my_model", ...)
state, history = train_from_config(cfg)
Source code in src/deqn_jax/models/__init__.py 
 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
def register_model(
    spec: ModelSpec,
    *,
    description: Optional[str] = None,
    overwrite: bool = False,
) -> None:
    """Register a model at runtime so ``load_model(spec.name)`` finds it.

    Equivalent to adding an entry to ``_MODELS`` and ``_DESCRIPTIONS``
    in this file, but at import / instantiation time. Intended for:

      * agent-codegen'd models in user projects (no fork required),
      * notebook prototyping,
      * plugin packages that add models to deqn-jax via ``register_model``
        on import.

    Args:
        spec: A populated ``ModelSpec``. Its ``spec.name`` field is the
            registry key. The framework consumes this object via the same
            ``load_model`` lookup as in-tree models.
        description: Human-readable one-liner for ``list_models``. Defaults
            to an empty string if omitted.
        overwrite: If True, an existing entry with the same name is
            replaced. If False (default) and ``spec.name`` is already
            registered, raises ``ValueError`` so you don't shadow an
            in-tree model by accident.

    Raises:
        ValueError: If ``spec.name`` is empty, or already registered and
            ``overwrite=False``.
        TypeError: If ``spec`` is not a ``ModelSpec``.

    Example::

        from deqn_jax.api import ModelSpec, register_model

        MY_MODEL = ModelSpec(name="my_model", ...)
        register_model(MY_MODEL, description="My agent-built model")

        # Now usable everywhere:
        from deqn_jax.api import load_model, train_from_config, TrainConfig
        cfg = TrainConfig(model="my_model", ...)
        state, history = train_from_config(cfg)
    """
    if not isinstance(spec, ModelSpec):
        raise TypeError(
            f"register_model expects a ModelSpec; got {type(spec).__name__}"
        )
    if not spec.name:
        raise ValueError("ModelSpec.name must be a non-empty string")
    if spec.name in _MODELS and not overwrite:
        raise ValueError(
            f"Model {spec.name!r} is already registered. Pass overwrite=True "
            f"to replace it, or pick a different name. Currently registered: "
            f"{sorted(_MODELS)}"
        )
    _MODELS[spec.name] = spec
    _DESCRIPTIONS[spec.name] = description or ""

unregister_model 

unregister_model(name: str) -> None

Remove a runtime-registered model. No-op if not present.

Mainly useful in tests so a leaked register_model from one test case doesn't bleed into another. Note: in-tree models can also be removed this way, but doing so is rarely a good idea.

Source code in src/deqn_jax/models/__init__.py 
137
138
139
140
141
142
143
144
145
def unregister_model(name: str) -> None:
    """Remove a runtime-registered model. No-op if not present.

    Mainly useful in tests so a leaked ``register_model`` from one test
    case doesn't bleed into another. Note: in-tree models can also be
    removed this way, but doing so is rarely a good idea.
    """
    _MODELS.pop(name, None)
    _DESCRIPTIONS.pop(name, None)

deqn_jax.models.variable_spec

Variable access helpers - named views over state/policy arrays.

Instead of fragile index slicing like state[:, 0], use:

s = unpack_state(state, MODEL)
s.k  # capital
s.z  # TFP

JAX traces through NamedTuples efficiently.

VariableSpec

Specification for model variables with named access.

Usage

spec = VariableSpec( state_names=("k", "z"), policy_names=("sav_rate",), )

Unpack

s = spec.unpack_state(state_array) p = spec.unpack_policy(policy_array)

Access by name

capital = s.k savings = p.sav_rate

Pack back

new_state = spec.pack_state(s._replace(k=new_k))

Source code in src/deqn_jax/models/variable_spec.py 
 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
127
128
129
130
131
132
133
134
class VariableSpec:
    """Specification for model variables with named access.

    Usage:
        spec = VariableSpec(
            state_names=("k", "z"),
            policy_names=("sav_rate",),
        )

        # Unpack
        s = spec.unpack_state(state_array)
        p = spec.unpack_policy(policy_array)

        # Access by name
        capital = s.k
        savings = p.sav_rate

        # Pack back
        new_state = spec.pack_state(s._replace(k=new_k))
    """

    def __init__(
        self,
        state_names: Tuple[str, ...],
        policy_names: Tuple[str, ...],
    ):
        self.state_names = state_names
        self.policy_names = policy_names
        self.n_states = len(state_names)
        self.n_policies = len(policy_names)

        # Create NamedTuple types
        self.StateType = make_state_type(state_names)
        self.PolicyType = make_policy_type(policy_names)

        # Create index lookup dicts
        self.state_idx = {name: i for i, name in enumerate(state_names)}
        self.policy_idx = {name: i for i, name in enumerate(policy_names)}

    # Return / parameter types are ``Any`` rather than ``NamedTuple``
    # because the actual tuple class is built at __init__ from
    # ``state_names``/``policy_names`` (see ``make_state_type``); a
    # static checker can't see those field names through the abstract
    # ``NamedTuple`` base. Using ``Any`` here lets ``s.pi_lag``,
    # ``p.sav_rate``, etc. type-check at every model's
    # ``equations.py`` call site instead of producing 200+ spurious
    # reportAttributeAccessIssue errors.
    def unpack_state(self, arr: Array) -> Any:
        """Unpack state array into named fields."""
        return unpack_array(arr, self.state_names, self.StateType)

    def unpack_policy(self, arr: Array) -> Any:
        """Unpack policy array into named fields."""
        return unpack_array(arr, self.policy_names, self.PolicyType)

    def pack_state(self, state: Any) -> Array:
        """Pack state NamedTuple back into array."""
        return pack_array(state)

    def pack_policy(self, policy: Any) -> Array:
        """Pack policy NamedTuple back into array."""
        return pack_array(policy)

    def get_state_idx(self, name: str) -> int:
        """Get index for state variable by name."""
        return self.state_idx[name]

    def get_policy_idx(self, name: str) -> int:
        """Get index for policy variable by name."""
        return self.policy_idx[name]

unpack_state 

unpack_state(arr: Array) -> Any

Unpack state array into named fields.

Source code in src/deqn_jax/models/variable_spec.py 
112
113
114
def unpack_state(self, arr: Array) -> Any:
    """Unpack state array into named fields."""
    return unpack_array(arr, self.state_names, self.StateType)

unpack_policy 

unpack_policy(arr: Array) -> Any

Unpack policy array into named fields.

Source code in src/deqn_jax/models/variable_spec.py 
116
117
118
def unpack_policy(self, arr: Array) -> Any:
    """Unpack policy array into named fields."""
    return unpack_array(arr, self.policy_names, self.PolicyType)

pack_state 

pack_state(state: Any) -> Array

Pack state NamedTuple back into array.

Source code in src/deqn_jax/models/variable_spec.py 
120
121
122
def pack_state(self, state: Any) -> Array:
    """Pack state NamedTuple back into array."""
    return pack_array(state)

pack_policy 

pack_policy(policy: Any) -> Array

Pack policy NamedTuple back into array.

Source code in src/deqn_jax/models/variable_spec.py 
124
125
126
def pack_policy(self, policy: Any) -> Array:
    """Pack policy NamedTuple back into array."""
    return pack_array(policy)

get_state_idx 

get_state_idx(name: str) -> int

Get index for state variable by name.

Source code in src/deqn_jax/models/variable_spec.py 
128
129
130
def get_state_idx(self, name: str) -> int:
    """Get index for state variable by name."""
    return self.state_idx[name]

get_policy_idx 

get_policy_idx(name: str) -> int

Get index for policy variable by name.

Source code in src/deqn_jax/models/variable_spec.py 
132
133
134
def get_policy_idx(self, name: str) -> int:
    """Get index for policy variable by name."""
    return self.policy_idx[name]

make_state_type 

make_state_type(names: Tuple[str, ...]) -> Type[NamedTuple]

Dynamically create a NamedTuple type for state variables.

Source code in src/deqn_jax/models/variable_spec.py 
18
19
20
21
22
23
def make_state_type(names: Tuple[str, ...]) -> Type[NamedTuple]:
    """Dynamically create a NamedTuple type for state variables."""
    # Functional-form NamedTuple call with a runtime-built fields list.
    # Both ty and pyright reject the non-literal fields argument; the
    # construction is only used to give attribute access at trace time.
    return NamedTuple("State", [(name, Array) for name in names])  # pyright: ignore[reportArgumentType]  # ty: ignore[invalid-named-tuple]

make_policy_type 

make_policy_type(
    names: Tuple[str, ...],
) -> Type[NamedTuple]

Dynamically create a NamedTuple type for policy variables.

Source code in src/deqn_jax/models/variable_spec.py 
26
27
28
def make_policy_type(names: Tuple[str, ...]) -> Type[NamedTuple]:
    """Dynamically create a NamedTuple type for policy variables."""
    return NamedTuple("Policy", [(name, Array) for name in names])  # pyright: ignore[reportArgumentType]  # ty: ignore[invalid-named-tuple]

unpack_array 

unpack_array(
    arr: Array, names: Tuple[str, ...], nt_type: Type
) -> NamedTuple

Unpack array columns into named fields.

Parameters:

Name Type Description Default
arr Array

Array of shape [batch, n_vars] or [n_vars]

 required
names Tuple[str, ...]

Tuple of variable names

 required
nt_type Type

NamedTuple type to create

 required

Returns:

Type Description
NamedTuple

NamedTuple with named fields
Source code in src/deqn_jax/models/variable_spec.py 
31
32
33
34
35
36
37
38
39
40
41
42
43
44
def unpack_array(arr: Array, names: Tuple[str, ...], nt_type: Type) -> NamedTuple:
    """Unpack array columns into named fields.

    Args:
        arr: Array of shape [batch, n_vars] or [n_vars]
        names: Tuple of variable names
        nt_type: NamedTuple type to create

    Returns:
        NamedTuple with named fields
    """
    if arr.ndim == 1:
        return nt_type(*[arr[i] for i in range(len(names))])
    return nt_type(*[arr[:, i] for i in range(len(names))])

pack_array 

pack_array(nt: NamedTuple) -> Array

Pack NamedTuple fields back into array.

Parameters:

Name Type Description Default
nt NamedTuple

NamedTuple with array fields

 required

Returns:

Type Description
Array

Array of shape [batch, n_vars] or [n_vars]
Source code in src/deqn_jax/models/variable_spec.py 
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
def pack_array(nt: NamedTuple) -> Array:
    """Pack NamedTuple fields back into array.

    Args:
        nt: NamedTuple with array fields

    Returns:
        Array of shape [batch, n_vars] or [n_vars]
    """
    # ``list(nt)`` is typed ``list[object]`` because ``NamedTuple`` (the
    # abstract base) isn't generic over its element type; the fields are
    # all Arrays at runtime.
    values: List[Array] = list(nt)
    if values[0].ndim == 0:
        return jnp.stack(values)
    return jnp.stack(values, axis=1)

make_init_state_fn 

make_init_state_fn(
    state_names: Tuple[str, ...],
    init_specs: Dict[str, Dict[str, Any]],
) -> Callable[..., Array]

Build an init_state_fn from per-variable distributional specs.

Parameters:

Name Type Description Default
state_names Tuple[str, ...]

Ordered tuple of state variable names.

 required
init_specs Dict[str, Dict[str, Any]]

Dict mapping state name to a spec {"distribution": <name>, "kwargs": {...}}. State names without an entry default to constant: 0.

 required

Returns:

Type Description
Callable[..., Array]

A function with signature (key, batch_size, constants) -> Array
Callable[..., Array]

producing a [batch_size, len(state_names)] initial state.
Callable[..., Array]

Constants are not consumed by default (distributions fix their
Callable[..., Array]

kwargs at spec time) but the signature matches ModelSpec's
Callable[..., Array]

init_state_fn contract so the framework can swap it in.

Unknown distribution names raise ValueError at build time.

Source code in src/deqn_jax/models/variable_spec.py 
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
def make_init_state_fn(
    state_names: Tuple[str, ...],
    init_specs: Dict[str, Dict[str, Any]],
) -> Callable[..., Array]:
    """Build an init_state_fn from per-variable distributional specs.

    Args:
        state_names: Ordered tuple of state variable names.
        init_specs: Dict mapping state name to a spec
            ``{"distribution": <name>, "kwargs": {...}}``. State names
            without an entry default to ``constant: 0``.

    Returns:
        A function with signature ``(key, batch_size, constants) -> Array``
        producing a ``[batch_size, len(state_names)]`` initial state.
        Constants are not consumed by default (distributions fix their
        kwargs at spec time) but the signature matches ModelSpec's
        ``init_state_fn`` contract so the framework can swap it in.

    Unknown distribution names raise ValueError at build time.
    """

    # Validate early so model-construction errors are obvious.
    for name, spec in init_specs.items():
        if name not in state_names:
            raise ValueError(
                f"init_specs contains unknown state '{name}'. "
                f"Known states: {state_names}"
            )
        dist = spec.get("distribution")
        if dist not in _DISTRIBUTION_SAMPLERS:
            raise ValueError(
                f"Unknown distribution '{dist}' for state '{name}'. "
                f"Available: {sorted(_DISTRIBUTION_SAMPLERS)}"
            )

    def init_state_fn(key, batch_size: int, constants: Dict[str, float]):
        keys = jax.random.split(key, len(state_names))
        columns = []
        for i, name in enumerate(state_names):
            if name in init_specs:
                spec = init_specs[name]
                sampler = _DISTRIBUTION_SAMPLERS[spec["distribution"]]
                kwargs = dict(spec.get("kwargs", {}))
                columns.append(sampler(keys[i], (batch_size,), **kwargs))
            else:
                columns.append(jnp.zeros((batch_size,)))
        return jnp.stack(columns, axis=1)

    return init_state_fn