Fix for jax>=0.4.7 (#1595)

* Fix for jax 0.4.11

* Require jax, jaxlib version >= 0.4.7

Author: fehiepsi

numpyro/ops/provenance.py

@@ -2,12 +2,11 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import jax
-from jax._src.pjit import pjit_p
 from jax.api_util import flatten_fun, shaped_abstractify
 import jax.core as core
+from jax.experimental.pjit import pjit_p
 from jax.interpreters.partial_eval import trace_to_jaxpr_dynamic
 from jax.interpreters.pxla import xla_pmap_p
-from jax.interpreters.xla import xla_call_p
 import jax.linear_util as lu
 import jax.numpy as jnp
 
@@ -102,7 +101,6 @@ def track_deps_call_rule(eqn, provenance_inputs):
 
 
 track_deps_rules[core.call_p] = track_deps_call_rule
-track_deps_rules[xla_call_p] = track_deps_call_rule
 track_deps_rules[xla_pmap_p] = track_deps_call_rule
 
 

setup.py

@@ -9,8 +9,8 @@
 from setuptools import find_packages, setup
 
 PROJECT_PATH = os.path.dirname(os.path.abspath(__file__))
-_jax_version_constraints = ">=0.4"
-_jaxlib_version_constraints = ">=0.4"
+_jax_version_constraints = ">=0.4.7"
+_jaxlib_version_constraints = ">=0.4.7"
 
 # Find version
 for line in open(os.path.join(PROJECT_PATH, "numpyro", "version.py")):