dsx.infer for filtering and smoothing without numpyro by theorashid · Pull Request #238 · BasisResearch/dynestyx

theorashid · 2026-05-26T22:42:03Z

This is for #213.

So far it is just for Filter and Smoother. Simulator will have to come.

Basic idea is that Filter/Smoother no longer call numpyro.factor inline, it happens in dsx.sample. test_filters and test_smoothers still pass, happy days.

# non-numpyro path: returns InferResult directly (no side effects)
with Filter(filter_config=KFConfig(...)):
    result = dsx.infer("f", dynamics, obs_times=t, obs_values=y)
loss = -result.marginal_loglik

# numpyro path (unchanged from before):
with Filter(filter_config=KFConfig(...)):
    dsx.sample("f", dynamics, obs_times=t, obs_values=y)

In the background, sample just calls infer and then registers sites. dsx.infer returns InferResult which carries marginal_loglik, states, dists, and a private _register_numpyro_sites callback.

def sample(...):
    result = infer(...)
    if isinstance(result, InferResult) and result._register_numpyro_sites is not None:
        result._register_numpyro_sites(name)
    return result

The main work can be seen in the test_*_standalone.py, which I based from my cuthbert-models repo. This is how I would see it being used.

def test_infer_optax_mle():
    """Use dsx.infer + optax to do MLE without numpyro."""
    obs_times, obs_values = _make_data()

    def neg_loglik(alpha):
        dynamics = _make_lti_dynamics(alpha)
        with Filter(filter_config=KFConfig(filter_source="cuthbert")):
            result = dsx.infer(
                "f", dynamics, obs_times=obs_times, obs_values=obs_values
            )
        return -result.marginal_loglik

    optimizer = optax.adam(1e-2)
    alpha = jnp.array(0.3)
    opt_state = optimizer.init(alpha)

    initial_loss = neg_loglik(alpha)
    grad_fn = jax.grad(neg_loglik)

    for _ in range(20):
        grads = grad_fn(alpha)
        updates, opt_state = optimizer.update(grads, opt_state)
        alpha = optax.apply_updates(alpha, updates)

    final_loss = neg_loglik(alpha)
    assert final_loss < initial_loss

how these changes were made

Of course, this was largely done by burning tokens and handholding so that it matches the design that I (and then with feedback from both of you) wanted. I am not as familiar with the internals of the library, so if there are other places beyond the diff that you think these changes might affect, let me know where to look – I was relying a bit on existing tests not breaking.

smaller design things

I largely tried to keep everything that was there before to reduce the size of the refactor. But:

I did make BaseLogFactorAdder (and the Filter equivalent) an ABC because I think it made sense.
Plate, as before, does not register per-field sites.
Maybe we should rename _sample_intp to _infer_intp.
InferResult a __call__ shim to satisfy FunctionOfTime protocol (needed because some model functions return dsx.sample(...)). Effectful uses the @defop return annotation to decide the type of fwd() returns. InferResult.__call__ raises NotImplementedError to satisfy the protocol — we can't change it without reworking Simulator.

Decompose dsx.sample into dsx.infer (pure computation, returns InferResult) + numpyro site registration (via callback). Filter and Smoother are now numpyro-free: they compute results and return InferResult with a deferred _register_numpyro_sites callback that dsx.sample fires. All integration backends return (marginal_loglik, states, dists) tuples with no side effects.

theorashid · 2026-05-26T23:23:26Z

On the testing suite, sometimes tests/test_hierarchical_simulator_discretizer_smokes.py, tests/test_science/test_discrete_time_l63_mcmc.py are flaky. tests/test_science/test_hmm.py::test_mcmc_inference breaks for TypeError: only 0-dimensional arrays can be converted to Python scalars. test_science is pretty slow in general

DanWaxman · 2026-05-27T21:12:29Z

Thanks Theo! This seems directionally about right, I like the implementation strategy!

In the background, sample just calls infer and then registers sites. dsx.infer returns InferResult which carries marginal_loglik, states, dists, and a private _register_numpyro_sites callback.

That makes sense! It will be a tiny bit tricky to get working with Simulators, I think, but that's probably okay. One tricky part is you're allowed to stack simulators with filters, i.e., with Simulator(), Filter(): dsx.sample(...); this allows one to sample from the filtering/posterior predictive. We'll need to pass the corresponding InferResult and append the necessary information.

I am not as familiar with the internals of the library, so if there are other places beyond the diff that you think these changes might affect, let me know where to look

I don't have any off the top of my head, besides the aforementioned interaction with Simulators. The other main test will be running all the notebooks in the documentation from scratch and making sure the results are qualitatively similar, but I think that can come in a bit.

I did make BaseLogFactorAdder (and the Filter equivalent) an ABC because I think it made sense.

Agreed, thanks!

Maybe we should rename _sample_intp to _infer_intp.

Sure, I more-or-less agree.

sometimes tests/test_hierarchical_simulator_discretizer_smokes.py, tests/test_science/test_discrete_time_l63_mcmc.py are flaky

I'm surprised the simulator/discretizer smokes are flaky, I haven't run into that before... but it's not super surprising to me that the discrete_time_l63_mcmc is flaky. I think those were written before we made EnKF the discrete-time default.

test_science is pretty slow in general

Right... I think the test_science suite has somewhat fallen into disuse, and our tests in general are a bit of a mess (though I think with decent coverage -- just disorganized). It's been sitting on the backlog for a bit. I wouldn't worry too much about particularly slow test_science tests as long as docs in the notebooks are looking okay.

At the risk of getting ahead of myself on a draft PR, I think it makes sense to set up a staging branch for this. Then, we can try to land this PR in the staging branch; worry about Simulators and its various interactions afterwards; then worry about the documentation lift that this implies.

theorashid · 2026-05-28T18:05:29Z

I think the test_science suite has somewhat fallen into disuse

Right, now I've seen the workflows I can see they just run the tests ignoring test_science.

I can give Simulator a go if you want it in this PR, but I'm wary of keeping the size small so it's easy for you to review – up to you.

DanWaxman · 2026-05-29T13:29:58Z

I can give Simulator a go if you want it in this PR, but I'm wary of keeping the size small so it's easy for you to review – up to you.

I think it makes sense to keep the PRs small, but also want to minimize the amount of half-working features on the upstream. So I've changed the base branch to dsx-infer-staging, where we can work bit-by-bit in implementing dsx.infer(...).

From that perspective, feel free to mark as ready to review whenever you feel it's ready and I'll take a close look :) thanks again!!

theorashid · 2026-05-29T19:56:28Z

Just renamed _sample_intp to _infer_intp.

I think this is a good point. Smaller change, tests passing, not tooooo many files to check over. Then I'll use any feedback before doing the Simulator refactor.

…rs-smoothers

mattlevine22 · 2026-05-29T23:21:25Z

Yeah, you can ignore test_science for now.
Did you say tests/test_hierarchical_simulator_discretizer_smokes.py is having issues? That one should be a solid test, but it looks like everything passed in (green checks on your previous commits).
I just ran tutorials 04 (discrete-time filter + simulator roll-out) and 06 (SDE filter + simulator roll-out), and both worked and look right locally, so that is a good sign!
I think I prefer dsx.condition, as I worry that infer will sound "all-powerful" to some users trying to do parameter estimation. Infer will sound a bit weird when we use it for Simulator rollouts, but it is not CRAZY to say that a simulator rollout is "conditioned", even if only on obs_times.

DanWaxman

This is pretty close on my end! A few things:

After talking with Matt, I agree that dsx.condition(...) is likely to be clearer than dsx.infer(...). condition is much closer to the statistical "work" that infer is actually doing.
Some inconsistencies in docstring should be fixed
The guards during registering results are on marginal_loglikelihood being None, but currently, empty filters return an increment of 0.0. It's not super clear what this means for registering sites as a result.
- Would also be nice to have a test for empty filters

DanWaxman · 2026-06-01T14:46:52Z

+    registering numpyro.factor / numpyro.deterministic if needed.
+
+    Returns:
+        tuple: (marginal_loglik, posterior, filtered_dists).


Should match the other return docstring block

https://github.com/theorashid/dynestyx/blob/536ead8c9a99b71ad6ab7dd26b4c02b4cab767cd/dynestyx/inference/integrations/cd_dynamax/continuous_filter.py#L212-L218

DanWaxman · 2026-06-01T14:48:23Z

+    Returns:
+        tuple: (marginal_loglik, posterior, smoothed_dists).


Same here (should match other more detailed docstring)

DanWaxman · 2026-06-01T17:33:25Z

    obs_len = int(obs_values.shape[0])
    if obs_len == 0:
-        return []
+        return jnp.array(0.0), None, []


Suggested change

return jnp.array(0.0), None, []

return None, None, []

I think the regestration guard is actually on the MLL:

https://github.com/theorashid/dynestyx/blob/536ead8c9a99b71ad6ab7dd26b4c02b4cab767cd/dynestyx/inference/filters.py#L348-L350

DanWaxman · 2026-06-01T17:34:27Z

    t1 = int(obs_values.shape[0])
    if t1 == 0:
-        return []
+        return jnp.array(0.0), None, []


Suggested change

return jnp.array(0.0), None, []

return None, None, []

As before, I think the guarding is on the MLL:

https://github.com/theorashid/dynestyx/blob/536ead8c9a99b71ad6ab7dd26b4c02b4cab767cd/dynestyx/inference/smoothers.py#L333-L335

theorashid · 2026-06-01T22:24:12Z

I've still kept InferResult. ConditionResult sounds weird, it sounds like the result of some if-else logic rather than what we do here.

DanWaxman · 2026-06-02T01:51:36Z

Thanks, will take another look tomorrow!

I've still kept InferResult. ConditionResult sounds weird, it sounds like the result of some if-else logic rather than what we do here.

Hmm, I agree... Maybe ConditioningResult? I don't think this has to be a blocking point.

mattlevine22 · 2026-06-02T01:54:54Z

Agree not to let it block...I'd probably do ConditionedResult fwiw

DanWaxman

I think this looks very reasonable. Certainly good enough to merge into a staging branch, for me. There is an actual limitation of this API right now (we don't have a great way to deal with non-CRN filters without invoking the numpyro.seed handler), though. That should be dealt with at some point.

DanWaxman · 2026-06-02T03:29:53Z

+        if config.crn_seed is not None:
+            key = config.crn_seed
+        else:
+            import warnings  # noqa: PLC0415
+
+            with warnings.catch_warnings():
+                warnings.simplefilter("ignore")
+                key = numpyro.prng_key()  # returns None outside seed handler



This also needn't be blocking right now, but this is a legitimate shortcoming; we don't really have a way to deal with non-CRN filters without using the numpyro seed handler. This is non-trivial to fix (though one option is to basically replicate the numpyro.prng_key() function/seed handler, which are not very complex and are a rather straightforward way to implement global random seeds).

theorashid · 2026-06-02T16:29:09Z

ConditionedResult sounds good. I changed it – start as we mean to go on

Fix ty type error: cast states to jax.Array for HMM filter sites

68de9bc

DanWaxman changed the base branch from main to dsx-infer-staging May 29, 2026 13:28

Rename _sample_intp to _infer_intp

abbcc82

theorashid marked this pull request as ready for review May 29, 2026 19:55

Merge branch 'dsx-infer-staging' into refactor/decouple-numpyro-filte…

536ead8

…rs-smoothers

DanWaxman requested review from DanWaxman and mattlevine22 May 29, 2026 21:29

DanWaxman requested changes Jun 1, 2026

View reviewed changes

Address code review: rename to dsx.condition, fix empty obs, docstrings

61ccd0e

DanWaxman self-requested a review June 2, 2026 03:25

DanWaxman approved these changes Jun 2, 2026

View reviewed changes

Rename InferResult to ConditionedResult

5b1e386

		Returns:
		tuple: (marginal_loglik, posterior, smoothed_dists).

Conversation

theorashid commented May 26, 2026

how these changes were made

smaller design things

Uh oh!

theorashid commented May 26, 2026

Uh oh!

DanWaxman commented May 27, 2026

Uh oh!

theorashid commented May 28, 2026

Uh oh!

DanWaxman commented May 29, 2026

Uh oh!

theorashid commented May 29, 2026

Uh oh!

mattlevine22 commented May 29, 2026

Uh oh!

DanWaxman left a comment

Choose a reason for hiding this comment

Uh oh!

DanWaxman Jun 1, 2026

Choose a reason for hiding this comment

Uh oh!

DanWaxman Jun 1, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

DanWaxman Jun 1, 2026

Choose a reason for hiding this comment

Uh oh!

DanWaxman Jun 1, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

theorashid commented Jun 1, 2026

Uh oh!

DanWaxman commented Jun 2, 2026

Uh oh!

mattlevine22 commented Jun 2, 2026

Uh oh!

DanWaxman left a comment

Choose a reason for hiding this comment

Uh oh!

DanWaxman Jun 2, 2026

Choose a reason for hiding this comment

Uh oh!

theorashid commented Jun 2, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants