Clean up global search and replace

patrick-kidger · Dec 18, 2024 · 65646f3 · 65646f3
1 parent 37b7c1e
commit 65646f3
Show file tree

Hide file tree

Showing 10 changed files with 0 additions and 10 deletions.
diff --git a/equinox/_ad.py b/equinox/_ad.py
@@ -18,7 +18,6 @@
 import jax
 import jax._src.traceback_util as traceback_util
 import jax.core
-import jax.extend
 import jax.interpreters.ad as ad
 import jax.numpy as jnp
 import jax.tree_util as jtu

diff --git a/equinox/_enum.py b/equinox/_enum.py
@@ -3,7 +3,6 @@
 
 import jax._src.traceback_util as traceback_util
 import jax.core
-import jax.extend
 import jax.numpy as jnp
 import numpy as np
 from jaxtyping import Array, ArrayLike, Bool, Int

diff --git a/equinox/_errors.py b/equinox/_errors.py
@@ -9,7 +9,6 @@
 import jax
 import jax._src.traceback_util as traceback_util
 import jax.core
-import jax.extend
 import jax.lax as lax
 import jax.numpy as jnp
 import jax.tree_util as jtu

diff --git a/equinox/_jit.py b/equinox/_jit.py
@@ -11,7 +11,6 @@
 import jax._src.dispatch
 import jax._src.traceback_util as traceback_util
 import jax.core
-import jax.extend
 import jax.errors
 import jax.numpy as jnp
 from jaxtyping import PyTree

diff --git a/equinox/_misc.py b/equinox/_misc.py
@@ -1,6 +1,5 @@
 import jax
 import jax.core
-import jax.extend
 import jax.numpy as jnp
 from jaxtyping import Array
 

diff --git a/equinox/_vmap_pmap.py b/equinox/_vmap_pmap.py
@@ -8,7 +8,6 @@
 import jax
 import jax._src.traceback_util as traceback_util
 import jax.core
-import jax.extend
 import jax.numpy as jnp
 import jax.tree_util as jtu
 import numpy as np

diff --git a/equinox/debug/_backward_nan.py b/equinox/debug/_backward_nan.py
@@ -3,7 +3,6 @@
 import jax
 import jax._src.traceback_util as traceback_util
 import jax.core
-import jax.extend
 import jax.numpy as jnp
 import jax.tree_util as jtu
 

diff --git a/equinox/debug/_dce.py b/equinox/debug/_dce.py
@@ -2,7 +2,6 @@
 
 import jax
 import jax.core
-import jax.extend
 import jax.numpy as jnp
 import jax.tree_util as jtu
 from jaxtyping import PyTree

diff --git a/equinox/internal/_loop/checkpointed.py b/equinox/internal/_loop/checkpointed.py
@@ -54,7 +54,6 @@
 
 import jax
 import jax.core
-import jax.extend
 import jax.lax as lax
 import jax.numpy as jnp
 import jax.tree_util as jtu

diff --git a/tests/test_tree.py b/tests/test_tree.py
@@ -3,7 +3,6 @@
 import equinox as eqx
 import jax
 import jax.core
-import jax.extend
 import jax.nn as jnn
 import jax.numpy as jnp
 import jax.random as jrandom