Update requirements and allow new versions of JAX.

PiperOrigin-RevId: 421345196

requirements/requirements.txt

@@ -1,6 +1,6 @@
 absl-py>=0.9.0
 chex>=0.0.8
 distrax>=0.0.2
-jax>=0.1.55,<=0.2.21  # remove once TFP is compatible with JAX>=0.2.22
+jax>=0.1.55
 jaxlib>=0.1.37
 numpy>=1.18.0