Update JAX version.

hbq1 · DistraxDev · commit fd0272024bd1 · 2021-05-18T03:29:25.000-07:00
PiperOrigin-RevId: 373741558
diff --git a/requirements/requirements.txt b/requirements/requirements.txt
@@ -1,7 +1,6 @@
 absl-py>=0.9.0
 chex>=0.0.5
-# Change the sign once TFP is fixed
-jax<=0.2.11
+jax>=0.2.12
 jaxlib>=0.1.37
 numpy>=1.18.0
-tensorflow-probability>=0.12.1
+tfp-nightly[jax]
