From 6bcef7564fb35efb440d973947c48816318cdac7 Mon Sep 17 00:00:00 2001 From: Matthew Levine Date: Tue, 1 Sep 2026 13:34:16 -0400 Subject: [PATCH] migrate dataclasses to equinox where possible for jax compat --- dynestyx/evaluation/handlers.py | 14 ++++++++---- dynestyx/inference/filters.py | 15 ++++++++++--- dynestyx/inference/observation_predictions.py | 5 ++--- dynestyx/inference/smoothers.py | 9 +++++--- dynestyx/observation_missingness.py | 13 +++++------ dynestyx/simulation/base.py | 17 +++++++++----- dynestyx/types.py | 22 +++++++------------ 7 files changed, 56 insertions(+), 39 deletions(-) diff --git a/dynestyx/evaluation/handlers.py b/dynestyx/evaluation/handlers.py index 6483d016..b7f797dc 100644 --- a/dynestyx/evaluation/handlers.py +++ b/dynestyx/evaluation/handlers.py @@ -78,7 +78,10 @@ def _sample_ds( scoring_config=self.observation_scoring_config, plate_shapes=plate_shapes, ) - filtered_result.evaluation_result = evaluation_result + filtered_result = dataclasses.replace( + filtered_result, + evaluation_result=evaluation_result, + ) forwarded_result = fwd( name, @@ -95,9 +98,12 @@ def _sample_ds( evaluation_result=evaluation_result, **kwargs, ) - evaluation_result._register_numpyro_sites = chain_numpyro_site_registrations( - evaluation_result._register_numpyro_sites, - getattr(forwarded_result, "_register_numpyro_sites", None), + evaluation_result = dataclasses.replace( + evaluation_result, + _register_numpyro_sites=chain_numpyro_site_registrations( + evaluation_result._register_numpyro_sites, + getattr(forwarded_result, "_register_numpyro_sites", None), + ), ) return evaluation_result diff --git a/dynestyx/inference/filters.py b/dynestyx/inference/filters.py index 1b2d3c08..6f411445 100644 --- a/dynestyx/inference/filters.py +++ b/dynestyx/inference/filters.py @@ -77,6 +77,7 @@ from dynestyx.models import DynamicalModel from dynestyx.types import ( ConditionedResult, + EvaluationResult, FunctionOfTime, chain_numpyro_site_registrations, ) @@ -144,9 +145,17 @@ def _sample_ds( ) forwarded_register = getattr(forwarded_result, "_register_numpyro_sites", None) - result._register_numpyro_sites = chain_numpyro_site_registrations( - result._register_numpyro_sites, - forwarded_register, + if isinstance(forwarded_result, EvaluationResult): + result = dataclasses.replace( + result, + evaluation_result=forwarded_result, + ) + result = dataclasses.replace( + result, + _register_numpyro_sites=chain_numpyro_site_registrations( + result._register_numpyro_sites, + forwarded_register, + ), ) return result diff --git a/dynestyx/inference/observation_predictions.py b/dynestyx/inference/observation_predictions.py index cefdf29a..a41d147f 100644 --- a/dynestyx/inference/observation_predictions.py +++ b/dynestyx/inference/observation_predictions.py @@ -6,9 +6,9 @@ from __future__ import annotations -import dataclasses from typing import Any +import equinox as eqx import jax import jax.numpy as jnp import numpyro @@ -39,8 +39,7 @@ ) -@dataclasses.dataclass(frozen=True) -class PredictedObservationOutputs: +class PredictedObservationOutputs(eqx.Module): """Canonical predicted-observation outputs for Dynestyx filters.""" mean: Float[Array, "*plate time observation_dim"] | None = None diff --git a/dynestyx/inference/smoothers.py b/dynestyx/inference/smoothers.py index 06f1f3c7..b146baba 100644 --- a/dynestyx/inference/smoothers.py +++ b/dynestyx/inference/smoothers.py @@ -180,9 +180,12 @@ def _sample_ds( ) forwarded_register = getattr(forwarded_result, "_register_numpyro_sites", None) - result._register_numpyro_sites = chain_numpyro_site_registrations( - result._register_numpyro_sites, - forwarded_register, + result = dataclasses.replace( + result, + _register_numpyro_sites=chain_numpyro_site_registrations( + result._register_numpyro_sites, + forwarded_register, + ), ) return result diff --git a/dynestyx/observation_missingness.py b/dynestyx/observation_missingness.py index 01ce4f66..5efb9d93 100644 --- a/dynestyx/observation_missingness.py +++ b/dynestyx/observation_missingness.py @@ -2,10 +2,10 @@ from __future__ import annotations -import dataclasses from collections.abc import Callable from typing import Literal +import equinox as eqx import jax.numpy as jnp import jax.scipy as jsp import numpy as np @@ -27,8 +27,7 @@ MissingObservationStrategy = Literal["auto", "marginalize", "augment", "error"] -@dataclasses.dataclass -class MissingObservationMetadata: +class MissingObservationMetadata(eqx.Module): """Describe the missing entries in one observation array. Flattened indices list missing entries by time and then by component. @@ -54,10 +53,10 @@ class MissingObservationMetadata: missing_obs_times: Real[Array, " n_missing_obs"] missing_obs_coordinate_indices: Int[Array, " n_missing_obs"] | None missing_flat_indices: Int[Array, " n_missing_obs"] - observation_shape: tuple[int, ...] - has_missing: bool - has_partial_missing: bool - has_fully_missing_rows: bool + observation_shape: tuple[int, ...] = eqx.field(static=True) + has_missing: bool = eqx.field(static=True) + has_partial_missing: bool = eqx.field(static=True) + has_fully_missing_rows: bool = eqx.field(static=True) def _concrete_observation_mask( diff --git a/dynestyx/simulation/base.py b/dynestyx/simulation/base.py index 1e8e7a60..6401c894 100644 --- a/dynestyx/simulation/base.py +++ b/dynestyx/simulation/base.py @@ -32,6 +32,7 @@ ) from dynestyx.types import ( ConditionedResult, + EvaluationResult, SimulatedResult, chain_numpyro_site_registrations, ) @@ -504,13 +505,19 @@ def _register_self(site_name: str) -> None: downstream_register = getattr( downstream_result, "_register_numpyro_sites", None ) + combined_register = chain_numpyro_site_registrations( + _register_self, + results._register_numpyro_sites, + downstream_register, + ) + if isinstance(downstream_result, EvaluationResult): + return dataclasses.replace( + downstream_result, + _register_numpyro_sites=combined_register, + ) return dataclasses.replace( results, - _register_numpyro_sites=chain_numpyro_site_registrations( - _register_self, - results._register_numpyro_sites, - downstream_register, - ), + _register_numpyro_sites=combined_register, ) def simulate( diff --git a/dynestyx/types.py b/dynestyx/types.py index 283cd738..520c9656 100644 --- a/dynestyx/types.py +++ b/dynestyx/types.py @@ -1,6 +1,5 @@ """Shared typing helpers for dynamical systems.""" -import dataclasses from collections.abc import Callable from typing import Protocol, runtime_checkable @@ -17,8 +16,7 @@ def __call__( raise NotImplementedError() -@dataclasses.dataclass -class EvaluationResult: +class EvaluationResult(eqx.Module): """Outputs computed by an evaluation handler. Evaluation handlers attach this object to the ``ConditionedResult`` they @@ -26,16 +24,13 @@ class EvaluationResult: side-effect free while ``dsx.sample`` can register the same outputs later. """ - observation_scores: dict[str, Real[Array, "..."]] = dataclasses.field( - default_factory=dict - ) - _register_numpyro_sites: Callable[[str], None] | None = dataclasses.field( - default=None, repr=False + observation_scores: dict[str, Real[Array, "..."]] = eqx.field(default_factory=dict) + _register_numpyro_sites: Callable[[str], None] | None = eqx.field( + default=None, repr=False, static=True ) -@dataclasses.dataclass -class ConditionedResult: +class ConditionedResult(eqx.Module): """Common base for results from the NumPyro-free conditioning primitive. ``dsx.condition`` returns this type under both ``Filter`` and ``Smoother``. @@ -55,8 +50,8 @@ class ConditionedResult: dists: list | None = None predicted_observations: object = None evaluation_result: EvaluationResult | None = None - _register_numpyro_sites: Callable[[str], None] | None = dataclasses.field( - default=None, repr=False + _register_numpyro_sites: Callable[[str], None] | None = eqx.field( + default=None, repr=False, static=True ) def __call__( @@ -68,8 +63,7 @@ def __call__( ) -@dataclasses.dataclass -class LatentStateResult: +class LatentStateResult(eqx.Module): """Result of latent-state construction / scoring without NumPyro side effects. Let ``z = state_path_params`` denote the free variables used to