Assertion fail tips?

[jguhlin]

Hey, congrats on the pub!

I'm getting the following error periodically, Any ideas or suggestions? I'll keep twiddling and update as I go. I have a discussion question open as well.

Cheers,

---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)
Cell In[7], [line 1](vscode-notebook-cell:?execution_count=7&line=1)
----> [1](vscode-notebook-cell:?execution_count=7&line=1) results_pop0 = phlash.fit(train_chroms_ts, mutation_rate=7.4117E-09, niter=10_000, overlap=220, test_data=val_chroms_ts) # generation time is 7

File ~/development/hoiho_working/.pixi/envs/default/lib/python3.12/site-packages/phlash/mcmc.py:284, in fit(***failed resolving arguments***)
    281     assert jnp.isfinite(x).all()
    282     return x
--> [284](https://vscode-remote+ssh-002dremote-002bwork.vscode-resource.vscode-cdn.net/home/joseph/development/hoiho_working/ARG/tsinfer/~/development/hoiho_working/.pixi/envs/default/lib/python3.12/site-packages/phlash/mcmc.py:284) state = jax.tree.map(f, state1)
    285 _particles = state.particles
    286 if test_data is not None and i % 10 == 0:

File ~/development/hoiho_working/.pixi/envs/default/lib/python3.12/site-packages/jax/_src/tree.py:155, in map(f, tree, is_leaf, *rest)
    115 def map(f: Callable[..., Any],
    116         tree: Any,
    117         *rest: Any,
    118         is_leaf: Callable[[Any], bool] | None = None) -> Any:
    119   """Maps a multi-input function over pytree args to produce a new pytree.
    120 
    121   Args:
   (...)
    153     - :func:`jax.tree.reduce`
    154   """
--> [155](https://vscode-remote+ssh-002dremote-002bwork.vscode-resource.vscode-cdn.net/home/joseph/development/hoiho_working/ARG/tsinfer/~/development/hoiho_working/.pixi/envs/default/lib/python3.12/site-packages/jax/_src/tree.py:155)   return tree_util.tree_map(f, tree, *rest, is_leaf=is_leaf)
...
    280 def f(x):
--> [281](https://vscode-remote+ssh-002dremote-002bwork.vscode-resource.vscode-cdn.net/home/joseph/development/hoiho_working/ARG/tsinfer/~/development/hoiho_working/.pixi/envs/default/lib/python3.12/site-packages/phlash/mcmc.py:281)     assert jnp.isfinite(x).all()
    282     return x