jax_triton: migrate from deprecated jax.linear_util to jax.extend.lin…

…ear_util PiperOrigin-RevId: 576994737
jax-ml · Oct 26, 2023 · 1f41ec2 · 1f41ec2
1 parent 15a479e
commit 1f41ec2
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/jax_triton/experimental/fusion/fusion.py b/jax_triton/experimental/fusion/fusion.py
@@ -21,7 +21,7 @@
 
 import jax
 from jax import lax
-from jax import linear_util as lu
+from jax.extend import linear_util as lu
 from jax.interpreters import partial_eval as pe
 from jax.interpreters import xla
 from jax._src import core

diff --git a/jax_triton/experimental/fusion/lowering.py b/jax_triton/experimental/fusion/lowering.py
@@ -20,11 +20,11 @@
 import jax
 from jax import api_util
 from jax import lax
-from jax import linear_util as lu
 from jax import tree_util
 from jax._src import core
 from jax._src import util
 from jax._src import source_info_util
+from jax.extend import linear_util as lu
 from jax.interpreters import partial_eval as pe
 
 from jax_triton.experimental.fusion import fusion