maint: Bump JAX version for numpy 2 compatibility

[yaugenst-flex]

Currently potentially leads to invalid env as per #1913 (comment)

[notion-workspace]

⬆️ Bump jax to avoid invalid env state

[daquinteroflex]

We just have to decide if we need to do a new release

[yaugenst-flex]

@momchil-flex adding you as a reviewer because of the release question. this is going into develop right now

[momchil-flex]

@dmarek-flex was the issue restricted to trying to use our jax plugin, or was it somehow manifesting even without that?

If restricted to the jax plugin we're probably fine... If anyone complains in the meantime we can tell them to upgrade jax.

[dmarek-flex]

@dmarek-flex was the issue restricted to trying to use our jax plugin, or was it somehow manifesting even without that?

For me the issue manifested when running the quickstart notebook, since optax uses jax.

[dmarek-flex]

Everything seems to be working as expected! In case it matters I found that on develop numpy>=2 will be installed in a fresh environment, while on pre/2.8 numpy <2 is installed. I think that is what happened to my environment, switching between branches and running the poetry install commands updated my numpy version.