Update to explicit jax version requirements for 0.11 branch

tensorflow · Sep 28, 2020 · db26e9a · db26e9a
1 parent a134096
commit db26e9a
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/setup.py b/setup.py
@@ -105,7 +105,7 @@ def has_ext_modules(self):
     ],
     keywords='tensorflow probability statistics bayesian machine learning',
     extras_require={  # e.g. `pip install tfp-nightly[jax]`
-        'jax': ['jax', 'jaxlib'],
+        'jax': ['jax==0.1.74', 'jaxlib==0.1.52'],
         'tfds': [TFDS_PACKAGE],
     }
 )
diff --git a/testing/install_test_dependencies.sh b/testing/install_test_dependencies.sh
@@ -150,7 +150,7 @@ install_python_packages() {
   python -m pip install $PIP_FLAGS $TF_NIGHTLY_PACKAGE==$TF_VERSION_STR
 
   # For the JAX backend.
-  python -m pip install jax jaxlib
+  python -m pip install jax==0.1.74 jaxlib==0.1.52
 
   # The following unofficial dependencies are used only by tests.
   # TODO(b/148685448): Unpin Hypothesis and coverage versions.