From f114e545ac23886cc6c5d54a77c73e36f090573a Mon Sep 17 00:00:00 2001
From: Brian Patton <22173987+brianwa84@users.noreply.github.com>
Date: Mon, 12 Apr 2021 09:50:46 -0400
Subject: [PATCH] Update setup.py

JAX 0.2.12 removed custom_transforms.
---
 setup.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/setup.py b/setup.py
index 7f57cd99a1..790eee53df 100644
--- 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],
     }
 )