Update local lax friedrichs for newer jax versions, fixes #2

StanfordASL / hj_reachability

Hamilton-Jacobi reachability analysis in JAX.

MIT License

103 stars 16 forks source link

Update local lax friedrichs for newer jax versions, fixes #2 #3

Closed a-doering closed 1 year ago

a-doering commented 1 year ago

[x] update code to work with newer jax versions, fixes #2 . Some of the used jax methods are depreceated and have been removed. jax.tree_util.tree_multimap has been removed since jax 0.3.16, jax.ops.index_update and jax.ops.index ave been removed since jax 0.3.2. Tested with jax==0.3.23 and jax_lib==0.3.22.
[ ] Specify Jax version in requirements.

schmrlng commented 1 year ago

Ah sorry, I had indeed fixed those JAX deprecations but never pushed the changes -- I'll tag a new version once I can get CI passing (looks like Python 3.6 is finally officially EOL).

Regarding specifying the JAX version requirement, I had been taking the approach of omitting it from requirements.txt (see, e.g., https://github.com/deepmind/dm-haiku#installation for an argument in favor of this which I had read at the time of package creation) but I'll try to do a quick survey of the JAX ecosystem soon to see what's standard these days.

schmrlng commented 1 year ago

It seems like including a JAX version requirement (with a caveat to install JAX first if accelerator support is desired) is reasonably standard these days, and in general it's probably good to have a known-working version number listed somewhere in the package.

Thanks for bumping these issues to my attention!