Model checkers

Validation and shape-inference helpers for DynamicalModel construction and checks.

Validation and shape-inference helpers for dynamical models.

_infer_observation_dim_in_plate_context(*, observation_model: Callable[[Real[Array, ' state_dim'] | Real[Array, ''], Real[Array, ' control_dim'] | Real[Array, ''] | None, Real[Array, '']], Any], x_probe: Real[Array, ' state_dim'] | Real[Array, ''], u_probe: Real[Array, ' control_dim'] | Real[Array, ''] | None, t_probe: Real[Array, ''], observation_dim: int | None) -> int

Infer observation dimension in plate context, falling back to explicit value.

_infer_vector_dim_from_distribution(distribution: Any, name: str, *, allow_batch_shape: bool = False) -> int

Infer scalar/vector event dimension from a NumPyro-compatible distribution.

When allow_batch_shape is true, leading batch dimensions are ignored and only the distribution event shape is used. This is needed inside dsx.plate where a vector-valued state distribution can have shape (plate_size, state_dim) but still represents a state_dim event.

Note: Name is used only for error messages.

_inside_numpyro_plate_context() -> bool

Return True when currently executing inside any active numpyro.plate frame.

_is_categorical_distribution(distribution: Any) -> bool

Return True for class-label categorical distributions.

This intentionally excludes one-hot categorical variants because the model logic here assumes scalar integer latent states.

_make_probe_state(initial_condition: Any, state_dim: int) -> Real[Array, ' state_dim'] | Real[Array, '']

Build a synthetic state value used for shape-check probes.

_unwrap_base_distribution(distribution: Any) -> Any

Peel common NumPyro wrapper distributions to inspect the base distribution.

NumPyro often wraps scalar/vector distributions in containers like Independent, ExpandedDistribution, or MaskedDistribution. For shape and categorical checks we want to reason about the base distribution semantics.

_validate_categorical_state(categorical_state: bool | None, inferred_categorical_state: bool) -> None

Ensure optional categorical_state agrees with inferred initial condition type.

_validate_continuous_state_evolution(state_evolution: Any, state_dim: int, x_probe: Real[Array, ' state_dim'] | Real[Array, ''], u_probe: Real[Array, ' control_dim'] | Real[Array, ''] | None, t_probe: Real[Array, '']) -> None

Validate the drift shape of a continuous-time state evolution.

_validate_continuous_time_flag(continuous_time: bool | None, inferred_continuous_time: bool) -> None

Ensure optional continuous_time agrees with inferred model type.

_validate_discrete_state_evolution_output_shape(state_evolution: Any, state_dim: int, x_probe: Real[Array, ' state_dim'] | Real[Array, ''], u_probe: Real[Array, ' control_dim'] | Real[Array, ''] | None, t_probe: Real[Array, '']) -> None

Validate a discrete-time state evolution against the inferred state dimension.

_validate_observation_dim(observation_dim: int | None, inferred_observation_dim: int) -> None

Ensure optional observation_dim agrees with inferred observation shape.

_validate_state_dim(state_dim: int | None, inferred_state_dim: int) -> None

Ensure optional state_dim agrees with inferred initial condition shape.