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

brianwa84 · web-flow · commit b1a08b3e0bf2 · 2021-04-14T15:18:25.000-04:00
Update setup.py
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],
     }
 )

Original file line number	Diff line number	Diff line change
`@@ -105,7 +105,7 @@ def has_ext_modules(self):`
`105`	`105`	`],`
`106`	`106`	`keywords='tensorflow probability statistics bayesian machine learning',`
`107`	`107`	extras_require={ # e.g. `pip install tfp-nightly[jax]`
`108`		`- 'jax': ['jax', 'jaxlib'],`
	`108`	`+ 'jax': ['jax<=0.2.11', 'jaxlib<=0.1.64'],`
`109`	`109`	`'tfds': [TFDS_PACKAGE],`
`110`	`110`	`}`
`111`	`111`	`)`