FIX: window_adaptation usage aligned to blackjax 1.2.3

[forgi86]

window_adaptation syntax has slightly changed and the current code does not work with blackjax 1.2.3

In bayesian-neural-network.ipynb, I replaced:

adapt = blackjax.window_adaptation(blackjax.nuts, potential, num_warmup)
final_state, kernel, _ = adapt.run(key_warmup, params)

with

adapt = blackjax.window_adaptation(blackjax.nuts, potential)
(state, parameters), _ = adapt.run(key_warmup, params, num_warmup)
kernel = blackjax.nuts(potential, **parameters).step

TODO: fix other examples accordingly

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[gerdm]

Hi @forgi86, could you please share which version of Jax and Blackjax did you use for the refactor? I'm getting a

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

I have

blackjax==1.2.3
jax==0.4.26
jaxlib==0.4.26
jaxopt==0.8.3

It seems that the latest version of jax does not have a max argument for jax.numpy.clip

Thanks!

[forgi86]

blackjax==1.2.3
jax==0.4.33
jaxid==0.2
jaxlib==0.4.33
jaxopt==0.8.3
jaxtyping==0.2.34