Replace jax.tree API with jax.tree_util for backwards compatibility

danielward27 / flowjax

https://danielward27.github.io/flowjax/

MIT License

82 stars 10 forks source link

Replace jax.tree API with jax.tree_util for backwards compatibility #143

Closed mdmould closed 6 months ago

mdmould commented 6 months ago

This replaces two instances of the new jax.tree API with the equivalent jax.tree_util functions. The former is only available since jax==0.4.25, so the other option would be to include >=0.4.25 in the dependencies.