This line is causing the issue. Lines where jnp.clip is used were recently updated due to #663 but currently the requirements allow the latest version of blackjax to be installed with versions of jax that do not recognise the new keyword argument. I encounter this with jax 0.4.26, provided with the Nvidia NGC jax docker image.
https://github.com/blackjax-devs/blackjax/blob/83bc3a04a64ac3379cd220932e88b44578f2e8e5/requirements-doc.txt#L7
This line is causing the issue. Lines where
jnp.clip
is used were recently updated due to #663 but currently the requirements allow the latest version of blackjax to be installed with versions of jax that do not recognise the new keyword argument. I encounter this with jax 0.4.26, provided with the Nvidia NGC jax docker image.