Skip to content

Adding jaxtyping and small shape improvements to inference #293

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
mattlevine22 wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Adding jaxtyping and small shape improvements to inference #293

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
9822a46
added jaxtyping to /inference/ and uniformly coerce obs_values/ctrl_v…
mattlevine22 Jul 31, 2026
50b510c
delay axis assurance until after batching-assessments
mattlevine22 Jul 31, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 5 additions & 5 deletions dynestyx/inference/configs/filter.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,8 +5,8 @@
import math
from typing import Literal

import jax
import jax.random as jr
from jaxtyping import PRNGKeyArray

ResamplingBaseMethod = Literal["systematic", "multinomial", "stratified"]
ResamplingDifferentiableMethod = Literal["stop_gradient", "straight_through", "soft"]
Expand Down Expand Up @@ -55,7 +55,7 @@ class BaseFilterConfig(abc.ABC):
this factor before the update. Values slightly above `1.0`
implement covariance inflation, which can improve robustness when
the model is misspecified. `None` disables rescaling.
crn_seed (jax.Array | None): Fix the PRNG key for stochastic filters
crn_seed (PRNGKeyArray | None): Fix the PRNG key for stochastic filters
(EnKF, PF). Useful when differentiating through the filter:
a fixed key makes the randomness a deterministic function of model
parameters. `None` draws a fresh key each call.
Expand All @@ -78,7 +78,7 @@ class BaseFilterConfig(abc.ABC):
record_max_elems: int = 100_000
filter_source: FilterSource | None = None
cov_rescaling: float | None = None
crn_seed: jax.Array | None = None
crn_seed: PRNGKeyArray | None = None


@dataclasses.dataclass
Expand Down Expand Up @@ -107,7 +107,7 @@ class EnKFConfig(BaseFilterConfig):
n_particles (int): Number of ensemble members. More members give a
better covariance estimate at higher compute cost. Defaults to
`30`.
crn_seed (jax.Array | None): Fixed PRNG key for the ensemble. Defaults
crn_seed (PRNGKeyArray | None): Fixed PRNG key for the ensemble. Defaults
to `jr.PRNGKey(0)`, i.e., common random numbers are used. This
can reduce variance in gradient-based learning, but introduces
further bias.
Expand Down Expand Up @@ -163,7 +163,7 @@ class EnKFConfig(BaseFilterConfig):
"""

n_particles: int = 30
crn_seed: jax.Array | None = dataclasses.field(
crn_seed: PRNGKeyArray | None = dataclasses.field(
default_factory=lambda: jr.PRNGKey(0)
)
perturb_measurements: bool | None = None
Expand Down
56 changes: 36 additions & 20 deletions dynestyx/inference/filters.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -74,7 +74,7 @@
FunctionOfTime,
chain_numpyro_site_registrations,
)
from dynestyx.utils import _dist_has_plate_batch_dims
from dynestyx.utils import _dist_has_plate_batch_dims, _ensure_trailing_event_axis

type SSMType = ContDiscreteNonlinearGaussianSSM | ContDiscreteNonlinearSSM

Expand All @@ -88,7 +88,7 @@ def _sample_ds(
name: str,
dynamics: DynamicalModel,
*,
plate_shapes=(),
plate_shapes: tuple[int, ...] = (),
obs_times: Real[Array, "*obs_time_plate obs_time"] | None = None,
obs_values: Real[Array, "*obs_value_plate obs_time observation_dim"]
| Real[Array, "*obs_value_plate obs_time"]
Expand Down Expand Up @@ -143,7 +143,7 @@ def _add_log_factors(
name: str,
dynamics: DynamicalModel,
*,
plate_shapes=(),
plate_shapes: tuple[int, ...] = (),
obs_times: Real[Array, "*obs_time_plate obs_time"] | None = None,
obs_values: Real[Array, "*obs_value_plate obs_time observation_dim"]
| Real[Array, "*obs_value_plate obs_time"]
Expand All @@ -161,7 +161,7 @@ def _build_infer_result(
) -> ConditionedResult: ...


def _default_filter_config(dynamics: DynamicalModel):
def _default_filter_config(dynamics: DynamicalModel) -> BaseFilterConfig:
"""Return appropriate default filter config when none specified."""
if dynamics.continuous_time:
return ContinuousTimeEnKFConfig()
Expand Down Expand Up @@ -237,7 +237,7 @@ def _add_log_factors(
name: str,
dynamics: DynamicalModel,
*,
plate_shapes=(),
plate_shapes: tuple[int, ...] = (),
obs_times: Real[Array, "*obs_time_plate obs_time"] | None = None,
obs_values: Real[Array, "*obs_value_plate obs_time observation_dim"]
| Real[Array, "*obs_value_plate obs_time"]
Expand Down Expand Up @@ -275,7 +275,6 @@ def _add_log_factors(
obs_values=obs_values,
mode="filter",
)

# Resolve PRNG key: use explicit seed from config, fall back to numpyro
# context (inside a seeded model), or None (deterministic filters don't need one).
if config.crn_seed is not None:
Expand Down Expand Up @@ -303,6 +302,11 @@ def _add_log_factors(
ctrl_values=ctrl_values,
)

if not isinstance(config, HMMConfigs):
obs_values = _ensure_trailing_event_axis(obs_values)
if ctrl_values is not None:
ctrl_values = _ensure_trailing_event_axis(ctrl_values)

if dynamics.continuous_time:
if not isinstance(config, ContinuousTimeConfigs):
valid = [c.__name__ for c in ContinuousTimeConfigs]
Expand Down Expand Up @@ -439,7 +443,7 @@ def _add_log_factors_batched(
)
output_kind = "continuous"

def compute_output(dyn, ot, ov, ovf, om, ct, cv, k):
def _compute_output(dyn, ot, ov, ovf, om, ct, cv, k):
return compute_continuous_filter(
dyn,
cast(ContinuousTimeFilterConfig, config),
Expand All @@ -454,7 +458,7 @@ def compute_output(dyn, ot, ov, ovf, om, ct, cv, k):
output_kind = "hmm"
uses_preprocessed_obs = True

def compute_output(dyn, ot, ov, ovf, om, ct, cv, k):
def _compute_output(dyn, ot, ov, ovf, om, ct, cv, k):
return compute_hmm_filter(
dyn,
obs_times=ot,
Expand All @@ -468,7 +472,7 @@ def compute_output(dyn, ot, ov, ovf, om, ct, cv, k):
if config.filter_source == "cuthbert":
output_kind = "cuthbert"

def compute_output(dyn, ot, ov, ovf, om, ct, cv, k):
def _compute_output(dyn, ot, ov, ovf, om, ct, cv, k):
return compute_cuthbert_filter(
dyn,
config,
Expand All @@ -482,7 +486,7 @@ def compute_output(dyn, ot, ov, ovf, om, ct, cv, k):
elif config.filter_source == "cd_dynamax":
output_kind = "cd_dynamax_discrete"

def compute_output(dyn, ot, ov, ovf, om, ct, cv, k):
def _compute_output(dyn, ot, ov, ovf, om, ct, cv, k):
return compute_cd_dynamax_discrete_filter(
dyn,
config,
Expand All @@ -499,6 +503,14 @@ def compute_output(dyn, ot, ov, ovf, om, ct, cv, k):
f"Unsupported filter config for plate: {type(config).__name__}"
)

def compute_output(dyn, ot, ov, ovf, om, ct, cv, k):
# Add scalar event axes after vmap removes plate dimensions.
if not isinstance(config, HMMConfigs):
ov = _ensure_trailing_event_axis(ov)
if cv is not None:
cv = _ensure_trailing_event_axis(cv)
return _compute_output(dyn, ot, ov, ovf, om, ct, cv, k)

# Pre-split keys for all plate members (needed for stochastic filters).
if key is not None:
# Ensure we use typed PRNG keys so split returns shape (total,)
Expand Down Expand Up @@ -669,13 +681,15 @@ def _filter_discrete_time(
key: PRNGKeyArray | None = None,
*,
obs_times: Real[Array, " obs_time"],
obs_values: Real[Array, "obs_time observation_dim"] | Real[Array, " obs_time"],
obs_values: Real[Array, "obs_time observation_dim"],
ctrl_times: Real[Array, " ctrl_time"] | None = None,
ctrl_values: Real[Array, "ctrl_time control_dim"]
| Real[Array, " ctrl_time"]
| None = None,
ctrl_values: Real[Array, "ctrl_time control_dim"] | None = None,
**kwargs,
) -> tuple[jax.Array | None, object | None, list[numpyro.distributions.Distribution]]:
) -> tuple[
Real[Array, ""] | None,
object | None,
list[numpyro.distributions.Distribution],
]:
"""Discrete-time marginal likelihood via cuthbert or cd-dynamax.

Filter type inferred from config class: KFConfig, EKFConfig, UKFConfig
Expand Down Expand Up @@ -725,13 +739,15 @@ def _filter_continuous_time(
key: PRNGKeyArray | None = None,
*,
obs_times: Real[Array, " obs_time"],
obs_values: Real[Array, "obs_time observation_dim"] | Real[Array, " obs_time"],
obs_values: Real[Array, "obs_time observation_dim"],
ctrl_times: Real[Array, " ctrl_time"] | None = None,
ctrl_values: Real[Array, "ctrl_time control_dim"]
| Real[Array, " ctrl_time"]
| None = None,
ctrl_values: Real[Array, "ctrl_time control_dim"] | None = None,
**kwargs,
) -> tuple[jax.Array, object, list[numpyro.distributions.Distribution]]:
) -> tuple[
Real[Array, ""],
object,
list[numpyro.distributions.Distribution],
]:
"""Continuous-time marginal likelihood via CD-Dynamax.

Supports: EnKF, DPF, EKF, UKF (inferred from config type).
Expand Down
45 changes: 24 additions & 21 deletions dynestyx/inference/integrations/cd_dynamax/continuous_filter.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,7 @@
"""Continuous-time filters via CD-Dynamax: KF, EnKF, DPF, EKF, UKF."""

from typing import Any

import jax
import jax.numpy as jnp
import numpyro.distributions as dist
Expand All @@ -11,6 +13,7 @@
from cd_dynamax.src.continuous_discrete_linear_gaussian_ssm.models import (
PosteriorGSSMFiltered,
)
from jaxtyping import Array, PRNGKeyArray, Real

from dynestyx.inference.configs.filter import (
ContinuousTimeDPFConfig,
Expand Down Expand Up @@ -39,12 +42,12 @@

def _config_to_cd_dynamax_filter_kwargs(
config: ContinuousTimeFilterConfig,
params,
obs_values,
obs_times,
ctrl_values,
key,
) -> dict:
params: Any,
obs_values: Real[Array, "obs_time observation_dim"],
obs_times: Real[Array, "obs_time one"],
ctrl_values: Real[Array, "ctrl_time control_dim"],
key: PRNGKeyArray | None,
) -> dict[str, Any]:
"""Build the filter_kwargs dict passed to cd_dynamax_model.filter()."""

# cd-dynamax uses the legacy PRNG key interface, but newer numpyro uses typed keys.
Expand Down Expand Up @@ -111,9 +114,9 @@ def _config_to_cd_dynamax_filter_kwargs(

def _run_linear_kf(
dynamics: DynamicalModel,
obs_times,
obs_values,
ctrl_values,
obs_times: Real[Array, "obs_time one"],
obs_values: Real[Array, "obs_time observation_dim"],
ctrl_values: Real[Array, "ctrl_time control_dim"],
filter_config: ContinuousTimeKFConfig,
) -> PosteriorGSSMFiltered:
"""Run exact continuous-discrete KF (AffineLinearDrift + constant diffusion + LinearGaussianObservation)."""
Expand All @@ -136,13 +139,13 @@ def _run_linear_kf(
def compute_continuous_filter(
dynamics: DynamicalModel,
filter_config: ContinuousTimeFilterConfig,
key: jax.Array | None = None,
key: PRNGKeyArray | None = None,
*,
obs_times: jax.Array,
obs_values: jax.Array,
ctrl_times=None,
ctrl_values=None,
):
obs_times: Real[Array, " obs_time"],
obs_values: Real[Array, "obs_time observation_dim"],
ctrl_times: Real[Array, " ctrl_time"] | None = None,
ctrl_values: Real[Array, "ctrl_time control_dim"] | None = None,
) -> Any:
"""Pure-JAX continuous-time filter computation (no numpyro side-effects)."""
obs_times_arr = jnp.asarray(obs_times)
if obs_times_arr.ndim == 1:
Expand Down Expand Up @@ -196,14 +199,14 @@ def run_continuous_filter(
name: str,
dynamics: DynamicalModel,
filter_config: ContinuousTimeFilterConfig,
key: jax.Array | None = None,
key: PRNGKeyArray | None = None,
*,
obs_times: jax.Array,
obs_values: jax.Array,
ctrl_times=None,
ctrl_values=None,
obs_times: Real[Array, " obs_time"],
obs_values: Real[Array, "obs_time observation_dim"],
ctrl_times: Real[Array, " ctrl_time"] | None = None,
ctrl_values: Real[Array, "ctrl_time control_dim"] | None = None,
**kwargs,
) -> tuple[jax.Array, object, list[dist.Distribution]]:
) -> tuple[Real[Array, ""], object, list[dist.Distribution]]:
"""Run continuous-time filter via CD-Dynamax.

Pure computation — no numpyro side-effects. Callers are responsible for
Expand Down
27 changes: 15 additions & 12 deletions dynestyx/inference/integrations/cd_dynamax/continuous_smoother.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,7 @@
"""Continuous-time smoothers via CD-Dynamax: KF, EKF."""

from typing import Any

import jax
import jax.numpy as jnp
import numpyro.distributions as dist
Expand All @@ -10,6 +12,7 @@
cdlgssm_smoother,
cdnlgssm_smoother,
)
from jaxtyping import Array, PRNGKeyArray, Real

from dynestyx.inference.configs.smoother import (
ContinuousTimeEKFSmootherConfig,
Expand All @@ -30,13 +33,13 @@
def compute_continuous_smoother(
dynamics: DynamicalModel,
smoother_config: ContinuousTimeSmootherConfig,
key: jax.Array | None = None,
key: PRNGKeyArray | None = None,
*,
obs_times: jax.Array,
obs_values: jax.Array,
ctrl_times=None,
ctrl_values=None,
):
obs_times: Real[Array, " obs_time"],
obs_values: Real[Array, "obs_time observation_dim"],
ctrl_times: Real[Array, " ctrl_time"] | None = None,
ctrl_values: Real[Array, "ctrl_time control_dim"] | None = None,
) -> Any:
"""Pure-JAX continuous-time smoother computation (no numpyro side-effects)."""
obs_times_arr = jnp.asarray(obs_times)
if obs_times_arr.ndim == 1:
Expand Down Expand Up @@ -118,14 +121,14 @@ def run_continuous_smoother(
name: str,
dynamics: DynamicalModel,
smoother_config: ContinuousTimeSmootherConfig,
key: jax.Array | None = None,
key: PRNGKeyArray | None = None,
*,
obs_times: jax.Array,
obs_values: jax.Array,
ctrl_times=None,
ctrl_values=None,
obs_times: Real[Array, " obs_time"],
obs_values: Real[Array, "obs_time observation_dim"],
ctrl_times: Real[Array, " ctrl_time"] | None = None,
ctrl_values: Real[Array, "ctrl_time control_dim"] | None = None,
**kwargs,
) -> tuple[jax.Array, object, list[dist.Distribution]]:
) -> tuple[Real[Array, ""], object, list[dist.Distribution]]:
"""Run continuous-time smoother via CD-Dynamax.

Pure computation — no numpyro side-effects. Callers are responsible for
Expand Down
Loading
Loading