Use a released version of optax rather than installing from Git. (#12173)

Use a released version of optax rather than installing from Git

examples/flax/language-modeling/requirements.txt

@@ -2,3 +2,4 @@ datasets >= 1.1.3
 jax>=0.2.8
 jaxlib>=0.1.59
 flax>=0.3.4
+optax>=0.0.8

examples/flax/text-classification/requirements.txt

@@ -2,4 +2,4 @@ datasets >= 1.1.3
 jax>=0.2.8
 jaxlib>=0.1.59
 flax>=0.3.4
-git+https://github.com/deepmind/optax.git
+optax>=0.0.8