TypeError: clip() got an unexpected keyword argument 'max'

[dirknbr]

This code ran fine some months ago but now fails, this was in a cpu colab

idata = bx_model.mcmc.blackjax_nuts(seed=jax.random.key(0), num_draws=1000, num_chains=2)

[/usr/local/lib/python3.10/dist-packages/jax/_src/linear_util.py](https://localhost:8080/#) in call_wrapped(self, *args, **kwargs)
    190 
    191     try:
--> 192       ans = self.f(*args, **dict(self.params, **kwargs))
    193     except:
    194       # Some transformations yield from inside context managers, so we have to

TypeError: clip() got an unexpected keyword argument 'max'

[ColCarroll]

I might need more details on versions of things and the entire stack trace, but as a guess it looks like this issue.

The fix would be to bump the version of jax you're running to past 0.4.27 (pip install -U "jax>=0.4.27", I think should do it)

[dirknbr]

thanks, that has indeed fixed the error