Tarmo Äijö

[This](https://jax.readthedocs.io/en/latest/_autosummary/jax.clear_caches.html) seems to solve the issue.

I'm not sure how actionable this is but I used the following code: ```python import numpy as np import numpyro, jax from numpyro import sample import numpyro.distributions as dist from...

@fehiepsi I did run a couple of experiments without NumPyro to see the effects of `jax.jit()`, `jax.clear_caches()`, and dynamic shapes. Did you mean something like that? These results suggest that...

This should work ```python print(potential_fn_gen()(init_params.z)) ``` Please check how`init_params.z` looks.

Please see [this](https://forum.pyro.ai/t/extended-kalman-filters/1243) discussion. As mentioned there, you'll probably find [stats_tools.py](https://github.com/standard-ai/recursive-bayesian-filtering/blob/master/stats_tools.py) and [demo_extended_kalman_filter.py](https://github.com/standard-ai/recursive-bayesian-filtering/blob/master/demos/demo_extended_kalman_filter.py) rather useful. Anyway, you need to use the properties of `EKFState` listed [here](https://docs.pyro.ai/en/stable/contrib.tracking.html#module-pyro.contrib.tracking.extended_kalman_filter). Indeed, this example is...