neural_tangents: migrate from deprecated jax.linear_util to jax.exten…

…d.linear_util PiperOrigin-RevId: 577220260
google · Nov 21, 2023 · 8d948dc · 8d948dc
1 parent 109796e
commit 8d948dc
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/neural_tangents/_src/empirical.py b/neural_tangents/_src/empirical.py
@@ -114,7 +114,6 @@
 from jax import jvp
 from jax import lax
 from jax import linear_transpose
-from jax import linear_util as lu
 from jax import vjp
 from jax import vmap
 from jax.core import Jaxpr
@@ -147,6 +146,12 @@
 from .utils.typing import VMapAxes
 from .utils.typing import VMapAxisTriple
 
+try:
+  # jax >=0.4.16
+  from jax.extend import linear_util as lu
+except ImportError:
+  from jax import linear_util as lu
+
 
 # LINEARIZATION AND TAYLOR EXPANSION