Fix JAX dependencies version constraints

- Update jax to >=0.3.25
- Update jaxlib to >=0.4.6 (earliest available for Python 3.10)
- Update flax to >=0.6.0 for compatibility

requirements.txt

@@ -12,9 +12,9 @@ seaborn
 gradio==3.50.0
 datasets==1.18.3
 wandb==0.12.11
-jax>=0.3.0,<0.4.0
-jaxlib>=0.3.0,<0.4.0
-flax>=0.4.0,<0.5.0
+jax>=0.3.25
+jaxlib>=0.4.6
+flax>=0.6.0
 optax
 sentencepiece==0.1.96
 aiohttp>=3.8.0,<3.9.0