Handle renaming of pjit_p to jit_p (#2052)

* Handle renaming of `pjit_p` to `jit_p`

* Replace tensorflow_probability with tfp-nightly

* Replace tensorflow_probability with tfp-nightly

* Update jaxns upper bound

* Bump jaxns version

* Use different jaxns version for Python 3.9

Author: tare

docs/requirements.txt

@@ -4,7 +4,8 @@ funsor
 ipython
 jax
 jaxlib
-jaxns==2.6.3
+jaxns==2.6.9;python_version>="3.10"
+jaxns==2.6.3;python_version<"3.10"
 Jinja2
 matplotlib
 multipledispatch
@@ -18,5 +19,5 @@ readthedocs-sphinx-search>=0.3.2
 sphinx>=5
 sphinx-gallery
 sphinx_rtd_theme
-tensorflow_probability
+tfp-nightly
 tqdm

numpyro/ops/provenance.py

@@ -3,8 +3,11 @@
 
 import jax
 from jax.api_util import flatten_fun, shaped_abstractify
-from jax.experimental.pjit import pjit_p
 
+try:
+    from jax.experimental.pjit import pjit_p
+except ImportError:
+    from jax.extend.core.primitives import jit_p as pjit_p
 try:
     import jax.extend.linear_util as lu
 except ImportError:

setup.py

@@ -64,14 +64,14 @@
             "flax",
             "funsor>=0.4.1",
             "graphviz",
-            "jaxns>=2.6.3,<=2.6.8",
+            "jaxns>=2.6.3,<=2.6.9",
             "matplotlib",
             "optax>=0.0.6",
             "pylab-sdk",  # jaxns dependency
             "pytest-cov",
             "pyyaml",  # flax dependency
             "requests",  # pylab dependency
-            "tensorflow_probability>=0.18.0",
+            "tfp-nightly",
         ],
         "examples": [
             "arviz",