LatentPathBuilder

One way to do Bayesian inference on a dynamical system is to "unroll" it and perform joint inference. For example, consider a dynamical system

\[ x_{t+1} \sim p(x_{t+1} \,|\, x_t) \]

with observations

\[ y_t \sim p(y_t \,|\, x_t). \]

Then we can perform inference on the joint distribution \(p(x_{0:T}, \theta \,|\, y_{1:T})\) by considering a probabilistic program that computes recursions in a for loop, conceptually something like

def model(ys):
  x_0 = numpyro.sample("x0", p_x0)
  theta = numpyro.sample("theta", p_theta)

  x_i = x_0

  for i in range(T):
    x_i = numpyro.sample(f"x_{i+1}", p_transition(x_i))
    numpyro.sample(f"y_{i+1}", p_observation(x_i), obs=ys[i])

LatentPathBuilder provides such a representation for joint learning, with support for arbitrary missingness in the data y. To support missingness and be more efficient, the pseudocode above is not how LatentPathBuilder is implemented, but is a useful mental model.

LatentPathBuilder is a NumPyro-facing handler: use it through dsx.sample(...) inside a NumPyro model. For pure-JAX trajectory scoring of fixed latent values, use dsx.log_prob(...) instead.

Conceptually:

• Simulator generates trajectories and observations.
• LatentPathBuilder constructs state_path_params, reconstructs the full state_path, and evaluates log p(x, y | ...).
• Filter and Smoother remain the preferred handlers when observations should be marginalized.

The main output names are intentionally distinct from simulator rollout names:

• f_state_path / f_state_path_times: explicit latent-path inference outputs from LatentPathBuilder
• f_states / f_times: rollout outputs from Simulator and dsx.simulate(...)

JIT and the builder cache

The shapes of NumPyro sites constructed by LatentPathBuilder depend on the missingness pattern. They must therefore be known at compile time.

There are three paths:

1. Run LatentPathBuilder concretely to compute the missingness.
2. Run it concretely once to fill its cache, and then run it under JIT. NumPyro MCMC routines typically do this automatically.
3. Provide the missingness metadata directly.

Path 2 must use the same LatentPathBuilder object because that object stores the cache. The recommended pattern is to define the model first and apply the handler when the model runs. Examples of the three paths follow.

1. Concrete execution

def conditioned_model(obs_times=None, obs_values=None):
    return dsx.sample(
        "f",
        dynamics,
        obs_times=obs_times,
        obs_values=obs_values,
    )

predictive = Predictive(conditioned_model, num_samples=10)
with dsx.LatentPathBuilder():
    result = predictive(
        prediction_key,
        obs_times=obs_times,
        obs_values=obs_values,
    )

2. Cached metadata

builder = dsx.LatentPathBuilder()
with builder:
    predictive(
        warmup_key,
        obs_times=obs_times,
        obs_values=obs_values,
    )
    result = jax.jit(predictive)(
        prediction_key,
        obs_times=obs_times,
        obs_values=obs_values,
    )

3. Explicit metadata

dsx.prepare_missing_observation_metadata(...) can create the metadata from concrete data. You can also create it directly:

obs_times = jnp.array([0.0, 1.0])
obs_values = jnp.array([[0.0, jnp.nan], [jnp.nan, 1.0]])

metadata = dsx.MissingObservationMetadata(
    missing_obs_times=jnp.array([0.0, 1.0]),
    missing_obs_coordinate_indices=jnp.array([1, 0], dtype=jnp.int32),
    missing_flat_indices=jnp.array([1, 2], dtype=jnp.int32),
    observation_shape=(2, 2),
    has_missing=True,
    has_partial_missing=True,
    has_fully_missing_rows=False,
)

All fields must match the layout of obs_values.

def conditioned_model(obs_times=None, obs_values=None):
    return dsx.sample(
        "f",
        dynamics,
        obs_times=obs_times,
        obs_values=obs_values,
        missing_obs_metadata=metadata,
    )

predictive = jax.jit(Predictive(conditioned_model, num_samples=10))
with dsx.LatentPathBuilder():
    result = predictive(
        prediction_key,
        obs_times=obs_times,
        obs_values=obs_values,
    )

Observation times and finite values can change. The missing-observation layout must agree with the cached or supplied layout. If you supply one metadata object, all plate members use its layout.

For plate members with different layouts, use one builder for all traced calls. Ragged LatentStateResult fields are flat lists of per-member arrays. Each NumPyro site keeps its shape, and the rightmost plate index varies fastest.

Implementation Details

To support arbitrary missingness, and for efficiency, the actual implementation of the LatentPathBuilder differs from the simple "unrolling" mental model. In particular, the entire state_path is intiialized as a numpyro site, via an improper uniform prior. The improper uniform prior is modifying so that calling its sample method (for example, as used in numpyro MCMC samplers by default) provides draws from the actual SSM prior. In particular:

• discrete state_path_params are drawn from the SSM prior;
• ODE state_path_params are drawn from the initial-condition prior and then solved on the requested time grid;
• remaining missing_obs_values are drawn from the observation model given the simulated state path.

After simulation, the model is scored according to the joint log-likelihood \(\log p(x, y | \theta)\), and registered as a numpyro site f_joint_log_prob_factor.

NumPyro-facing explicit latent-path inference.

LatentPathBuilder

Bases: ObjectInterpretation, HandlesSelf

Construct and score explicit latent state paths in a NumPyro model.

Use this handler as a context manager around dsx.sample(...). The builder creates array-valued sample sites, reconstructs the complete state path, and adds the joint state-observation log density to the NumPyro model.

Several different strategies are provided for missing data, specified by the missing_observation_strategy parameter:

• "marginalize" evaluates the observation density using only the observed components. The observation distribution must support this calculation (namely, a MultivariateNormal or IndependentDistribution).
• "augment" creates a "{name}_missing_obs_values" sample site, fills the missing components, and evaluates the completed observation. This method requires a continuous observation distribution.
• "auto" uses marginalization when supported and otherwise uses augmentation for a continuous observation distribution.
• "error" rejects partially observed vectors.

DiracIdentityObservation uses a separate path. Its missing components are inferred through state_path_params, and missing data support only "auto" or "augment".

A missingness layout can determine NumPyro site shapes. The builder infers and caches this layout whenever observations are concrete. Before an outer jax.jit passes obs_values dynamically, either reuse a builder that has seen concrete observations or pass eagerly prepared missing_obs_metadata to dsx.sample(...). JAX itself does not run the function eagerly before tracing.

Attributes:

Name	Type	Description
ode_simulator_config		ODE solver and integration settings used during deterministic continuous-time path reconstruction.
missing_observation_strategy		Method used to handle missing entries in obs_values.
chunk_size		Batch size passed to jax.lax.map while scoring transition and observation terms. The default, 0, evaluates all terms with one jax.vmap. None maps one term at a time. A positive integer evaluates batches of that size with jax.vmap.

Examples:

>>> builder = dsx.LatentPathBuilder()
>>> with builder:
...     result = dsx.sample(
...         "f",
...         dynamics,
...         obs_times=obs_times,
...         obs_values=obs_values,
...     )

__init__(ode_simulator_config: ODESimulatorConfig | None = None, missing_observation_strategy: MissingObservationStrategy = 'auto', chunk_size: int | None = 0) -> None

Initialize explicit latent-path inference.

Parameters:

Name	Type	Description	Default
ode_simulator_config	ODESimulatorConfig | None	ODE solver and integration settings used during deterministic continuous-time path reconstruction. Defaults to ODESimulatorConfig() when omitted.	None
missing_observation_strategy	MissingObservationStrategy	Method used to handle missing entries in obs_values, as described in the class documentation.	'auto'
chunk_size	int | None	Batch size passed to jax.lax.map while scoring transition and observation terms. The default, 0, evaluates all terms with one jax.vmap. None maps one term at a time. A positive integer evaluates batches of that size with jax.vmap.	0