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

* avoid breakage in old jax version without jax.extend

* fix lint

Author: fehiepsi

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

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