Relax jax/jaxlib version constraints

[samuela]

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

treeo/pyproject.toml

Lines 16 to 17 in a402f3f

	jax = "^0.2.24"
	jaxlib = "^0.1.73"

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]

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]

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

[cgarciae]

Hey! Tell me if #9 works for you.

[samuela]

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