Merge pull request tensorflow#1300 from tensorflow/brianwa84-patch-3

Update setup.py
NeilGirdhar · Apr 14, 2021 · b1a08b3 · b1a08b3
2 parents 43a9d6c + f114e54
commit b1a08b3
Showing 1 changed file with 1 addition and 1 deletion.
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.2.11', 'jaxlib<=0.1.64'],
         'tfds': [TFDS_PACKAGE],
     }
 )