blackjax-devs / blackjax

BlackJAX is a Bayesian Inference library designed for ease of use, speed and modularity.
https://blackjax-devs.github.io/blackjax/
Apache License 2.0
806 stars 105 forks source link

Current requirements allows for jax >=0.4.25, but new jnp.clip argument requires jax 0.4.27 #691

Open juehang opened 4 months ago

juehang commented 4 months ago

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.