Migrate jax.ShapedArray -> jax.core.ShapedArray

jax.ShapedArray is deprecated as of google/jax#15263: most users will never need to use ShapedArray directly, and so having it exposed in the top-level public namespace causes undue confusion. PiperOrigin-RevId: 521880270
google · Apr 19, 2023 · 396ad1d · 396ad1d
1 parent 31bc793
commit 396ad1d
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/neural_tangents/_src/stax/linear.py b/neural_tangents/_src/stax/linear.py
@@ -26,7 +26,8 @@
 from jax import numpy as np
 from jax import ops
 from jax import random
-from jax import ShapeDtypeStruct, ShapedArray, eval_shape, vmap
+from jax import ShapeDtypeStruct, eval_shape, vmap
+from jax.core import ShapedArray
 import jax.example_libraries.stax as ostax
 import numpy as onp
 from .requirements import Bool, Diagonal, get_diagonal_outer_prods, layer, mean_and_var, requires, supports_masking

diff --git a/neural_tangents/_src/stax/requirements.py b/neural_tangents/_src/stax/requirements.py
@@ -22,7 +22,8 @@
 import jax
 from jax import lax
 from jax import numpy as np
-from jax import eval_shape, ShapedArray
+from jax import eval_shape
+from jax.core import ShapedArray
 from jax.tree_util import tree_map, tree_all
 from ..utils import utils
 import dataclasses