avoid breakage in old jax version without jax.extend (#1647)

* avoid breakage in old jax version without jax.extend * fix lint
pyro-ppl · Sep 22, 2023 · 59a188d · 59a188d
1 parent 6e3f007
commit 59a188d
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 2 deletions.
diff --git a/numpyro/ops/provenance.py b/numpyro/ops/provenance.py
@@ -5,7 +5,11 @@
 from jax.api_util import flatten_fun, shaped_abstractify
 import jax.core as core
 from jax.experimental.pjit import pjit_p
-import jax.extend.linear_util as lu
+
+try:
+    import jax.extend.linear_util as lu
+except ImportError:
+    import jax.linear_util as lu
 from jax.interpreters.partial_eval import trace_to_jaxpr_dynamic
 from jax.interpreters.pxla import xla_pmap_p
 import jax.numpy as jnp

diff --git a/test/ops/test_provenance.py b/test/ops/test_provenance.py
@@ -8,7 +8,11 @@
 import jax
 from jax.api_util import flatten_fun_nokwargs
 import jax.core as core
-import jax.extend.linear_util as lu
+
+try:
+    import jax.extend.linear_util as lu
+except ImportError:
+    import jax.linear_util as lu
 import jax.numpy as jnp
 
 from numpyro.ops.provenance import eval_provenance