Consider using `numpyro.contrib.control_flow.scan` over `jax.lax.scan`

CDCgov / multisignal-epi-inference

Python package for statistical inference and forecast of epi models using multiple signals

https://cdcgov.github.io/multisignal-epi-inference/

9 stars 1 forks source link

Consider using `numpyro.contrib.control_flow.scan` over `jax.lax.scan` #213

Open damonbayer opened 5 days ago

damonbayer commented 5 days ago

Our use of jax.lax.scan has been fine so far, but the built-in numpyro.contrib.control_flow.scan offers additional features/compatibility with numpyro.