search

cgarciae / treeo

A small library for creating and manipulating custom JAX Pytree classes
https://cgarciae.github.io/treeo
MIT License
58 stars 4 forks source link

Relax jax/jaxlib version constraints #7

Open samuela opened 2 years ago

samuela commented 2 years ago

Now that jax 0.3.0 and jaxlib 0.3.0 have been released the version constraints in pyproject.toml are outdated.

https://github.com/cgarciae/treeo/blob/a402f3f69557840cfbee4d7804964b8e2c47e3f7/pyproject.toml#L16-L17

This corresponds to the version constraint jax<0.3.0,>=0.2.18 (https://python-poetry.org/docs/dependency-specification/#caret-requirements). Now that jax v0.3.0 has been released (https://github.com/google/jax/releases/tag/jax-v0.3.0) this doesn't work with the latest version. I think the same applies to jaxlib as well, since it also got upgraded to v0.3.0 (https://github.com/google/jax/releases/tag/jaxlib-v0.3.0).

cgarciae commented 2 years ago

Thanks @samuela ! Indeed we should upgrade to 0.3.0, to do this its just a mater of running:

poetry add --dev jax@latest jaxlib@latest

Upside is that since they were purposely marked as dev dependencies they won't block any user from upgrading jax on their project.

cgarciae commented 2 years ago

@samuela We could also just use >= instead of ^.

cgarciae commented 2 years ago

Hey! Tell me if #9 works for you.

samuela commented 2 years ago

Hey @cgarciae, thanks for putting #9 together! It does the trick for me!