Trajectory scoring

jax_morph.trajectory_logp

trajectory_logp(
    model: Model,
    history: BaseState,
    dt: float | Array = 1.0,
    *,
    score: str | Iterable[int] | None = None,
    checkpoint: bool = False,
) -> jaxlib._jax.Array

Score a recorded trace trajectory with a live reconstructed state carry.

history is the stacked complete state sequence s_0, ..., s_T returned by simulate(..., history=True). Frame zero is the detached initial conditioning state; frames one onward are detached recorded transition data. Deterministic computations and the replay semantics selected by each stochastic step remain live, including across numerical macro-step boundaries. score=None selects stochastic steps whose score_by_default is true, score='all' selects every stochastic step, and an explicit iterable selects model-step indices. checkpoint=True rematerializes each replayed macro-step during the backward pass.

Parameters:

• model (Model) –

  Model whose stochastic steps produced the history.

• history (BaseState) –

  Complete stacked state history s_0, ..., s_T.

• dt (float | Array, default: 1.0 ) –

  Macro-step size used to produce the history. Defaults to 1.0.

• score (str | Iterable[int] | None, default: None ) –

  Stochastic-step selection: None, 'all', or an iterable of model-step indices. Defaults to None (score the stochastic steps whose score_by_default is true).

• checkpoint (bool, default: False ) –

  Whether to rematerialize replayed macro-steps during backpropagation. Defaults to False.

Returns:

• Array –

  Per-macro-step log-probability terms with shape (T,). This is not the scalar joint

• Array –

  log-probability; call terms.sum() when a joint score is required.

---

jax_morph.transition_logp

transition_logp(
    model: Model,
    observed_state: BaseState,
    observed_next: BaseState,
    dt: float | Array = 1.0,
    *,
    score: str | Iterable[int] | None = None,
) -> jaxlib._jax.Array

Score one recorded transition conditioned on an explicitly observed state.

observed_state is detached as a statistical conditioning boundary. Stochastic traces are read from the detached observed_next state and replayed with parameters live. score has the same selection semantics as trajectory_logp.

Parameters:

• model (Model) –

  Model whose stochastic steps produced the transition.

• observed_state (BaseState) –

  State treated as the detached conditioning boundary.

• observed_next (BaseState) –

  Recorded next state containing the stochastic trace.

• dt (float | Array, default: 1.0 ) –

  Macro-step size used to produce the transition. Defaults to 1.0.

• score (str | Iterable[int] | None, default: None ) –

  Stochastic-step selection: None, 'all', or an iterable of model-step indices. Defaults to None (score the stochastic steps whose score_by_default is true).

Returns:

• Array –

  Scalar sum of the selected stochastic-step log-probabilities.