Skip to content

Staging Branch for Numpyro-Free Usage#255

Draft
DanWaxman wants to merge 2 commits into
mainfrom
dsx-infer-staging
Draft

Staging Branch for Numpyro-Free Usage#255
DanWaxman wants to merge 2 commits into
mainfrom
dsx-infer-staging

Conversation

@DanWaxman

@DanWaxman DanWaxman commented Jun 12, 2026

Copy link
Copy Markdown
Collaborator

Staging branch for numpyro-free usage.

* dsx.infer for numpyro-free filtering and smoothing

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.

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

* Rename _sample_intp to _infer_intp

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

* Rename InferResult to ConditionedResult

* Add NumPyro-free counterparts of tutorials 03 (filtering MLL) and 05 (optax MLE)

---------

Co-authored-by: Dan Waxman <dan.waxman1@gmail.com>
@DanWaxman

Copy link
Copy Markdown
Collaborator Author

@theorashid I set up this draft PR as staging for the numpyro-free stuff. My proposed workflow would be to open new PRs targeting this branch (dsx-infer-staging) for the remaining items (namely, support for simulators and the corresponding documentation changes).

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants