Skip to content

Model

jax_morph.Model #

Model(steps: Iterable[SimulationStep])

An ordered pipeline of steps with a validated field dataflow.

Model is callable: model(state, *, dt, key) advances the state by one macro-step. The forward pass is a pure sampler. Top-level trajectory_logp and transition_logp score its recorded stochastic choices through private replay machinery. Sampling and replay share one macro-step skeleton (_run).

Attributes:

  • steps –

    Ordered tuple of SimulationStep instances.

Methods:

  • __call__ –

    Advance one hybrid macro-step.

  • state_requires –

    Return the full step-declared state schema.

Parameters:

  • steps (Iterable[SimulationStep]) –

    Simulation steps in pipeline order.

Raises:

  • ValueError –

    If declared fields conflict or a stochastic trace is invalid.

state_requires #

state_requires() -> tuple[
    jax_morph.core.state.StateFieldSpec, ...
]

The union of every step's required field specs (deduplicated tuple of specs).

Deduplicates specs that are exactly equal and keeps every distinct field; raises when two steps declare the same field name with different properties (shape, dtype, scope, ...). The always-present base fields are the state's concern, not the model's; build_state_from_model combines them with this to form the full allocation schema.

Returns:

  • tuple[StateFieldSpec, ...] –

    A deduplicated tuple of StateFieldSpec objects required by the model's steps.

Raises:

  • ValueError –

    If two steps declare the same field name with different properties.