Relax jax/jaxlib version constraints about treeo

[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!