MAINT Clean up leftover Array = Any aliases in jax/_src/**.py

This seemingly minor change uncovered a few inconsistencies in existing code.
I silenced the new type errors via ``type: ignore` and added a TODO to
investigate later.

jax/_src/interpreters/batching.py

@@ -33,12 +33,12 @@
 from jax._src.interpreters import partial_eval as pe
 from jax._src.tree_util import (tree_unflatten, tree_flatten,
                                 register_pytree_node)
+from jax._src.typing import Array
 from jax._src.util import (unzip2, unzip3, safe_map, safe_zip, split_list,
                            canonicalize_axis, moveaxis, as_hashable_function,
                            curry, memoize, weakref_lru_cache)
 
 
-Array = Any
 map, unsafe_map = safe_map, map
 zip, unsafe_zip = safe_zip, zip
 
@@ -172,7 +172,8 @@ def shape_as_bdim(
   # This assumes that there is only one binder in the data_shape.
   ragged_axes = [(i, size.lengths) for i, size in enumerate(data_shape)
                  if isinstance(size, IndexedAxisSize)]
-  return make_batch_axis(len(data_shape), stacked_axis, ragged_axes)
+  # TODO(slebedev): This is a genuine type error, Array | T is not <: Array.
+  return make_batch_axis(len(data_shape), stacked_axis, ragged_axes)  # type: ignore[arg-type]
 
 
 def _update_annotation(
@@ -242,7 +243,8 @@ def to_elt(trace: Trace, get_idx: GetIdx, x: Vmappable, spec: MapSpec) -> Elt:
     (d, ias), = ((i, sz)  # type: ignore
                  for i, sz in enumerate(x.aval.elt_ty.shape)
                  if type(sz) is IndexedAxisSize)
-    batch_axis = make_batch_axis(x.data.ndim, 0, [(d+1, ias.lengths)])
+    # TODO(slebedev): This is a genuine type error, Array | T is not <: Array.
+    batch_axis = make_batch_axis(x.data.ndim, 0, [(d+1, ias.lengths)])  # type: ignore[list-item]
     return BatchTracer(trace, x.data, batch_axis)  # type: ignore
   elif isinstance(spec, int) or spec is None:
     spec = spec and canonicalize_axis(spec, len(np.shape(x)))
@@ -695,7 +697,8 @@ def canonicalize_segment_lengths(d: RaggedAxis) -> RaggedAxis:
           id(core.get_referent(segment_lengths)),
           (segment_lengths, pe.DBIdx(len(axis_map))))
       new_ragged_axes.append((ragged_axis, dbidx))
-    return RaggedAxis(d.stacked_axis, tuple(new_ragged_axes))
+      # TODO(slebedev): This is a geniune type error, pe.DBIdx is not <: Array.
+    return RaggedAxis(d.stacked_axis, tuple(new_ragged_axes))  # type: ignore[arg-type]
   new_dims = [canonicalize_segment_lengths(d)
               if isinstance(d, RaggedAxis) else d for d in dims]
   segment_lens = [s for s, _ in axis_map.values()]

jax/_src/lax/ann.py

@@ -70,7 +70,6 @@ def pmap_mips(qy, db, db_offset, db_size, k, recall_target):
 """
 
 from functools import partial
-from typing import Any
 
 import numpy as np
 
@@ -88,9 +87,7 @@ def pmap_mips(qy, db, db_offset, db_size, k, recall_target):
 from jax._src.lib.mlir import ir
 from jax._src.lib.mlir.dialects import func
 from jax._src.lib.mlir.dialects import hlo
-
-
-Array = Any
+from jax._src.typing import Array
 
 
 def approx_max_k(operand: Array,

jax/_src/lax/control_flow/for_loop.py

@@ -39,6 +39,7 @@
 from jax._src.state import primitives as state_primitives
 from jax._src.state import utils as state_utils
 from jax._src.state import types as state_types
+from jax._src.typing import Array
 from jax._src.util import (partition_list, merge_lists, safe_map, safe_zip,
                            split_list, split_dict)
 from jax._src.lax.control_flow import loops
@@ -53,7 +54,6 @@
 S = TypeVar('S')
 T = TypeVar('T')
 class Ref(Generic[T]): pass
-Array = Any
 
 ref_set = state_primitives.ref_set
 ref_get = state_primitives.ref_get

jax/_src/lax/control_flow/loops.py

@@ -52,6 +52,7 @@
 from jax._src.state import discharge as state_discharge
 from jax._src.numpy.ufuncs import logaddexp
 from jax._src.traceback_util import api_boundary
+from jax._src.typing import Array
 from jax._src.util import (partition_list, safe_map, safe_zip, split_list,
                            unzip2, weakref_lru_cache, merge_lists)
 import numpy as np
@@ -64,7 +65,6 @@
 zip = safe_zip
 
 T = TypeVar('T')
-Array = Any
 BooleanNumeric = Any  # A bool, or a Boolean array.
 
 ### Helper functions

jax/_src/lax/convolution.py

@@ -12,11 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import builtins
 from collections.abc import Sequence
 from functools import partial
 import operator
-from typing import Any, NamedTuple, Optional, Union
+from typing import NamedTuple, Optional, Union
 
 import numpy as np
 
@@ -28,14 +27,9 @@
 from jax._src.interpreters import mlir
 from jax._src.lax import lax
 from jax._src.lib.mlir.dialects import hlo
+from jax._src.typing import Array, DTypeLike
 
 
-_max = builtins.max
-
-Array = Any
-DType = Any
-Shape = core.Shape
-
 class ConvDimensionNumbers(NamedTuple):
   """Describes batch, spatial, and feature dimensions of a convolution.
 
@@ -62,7 +56,7 @@ def conv_general_dilated(
   dimension_numbers: ConvGeneralDilatedDimensionNumbers  = None,
   feature_group_count: int = 1, batch_group_count: int = 1,
   precision: lax.PrecisionLike = None,
-  preferred_element_type: Optional[DType] = None) -> Array:
+  preferred_element_type: Optional[DTypeLike] = None) -> Array:
   """General n-dimensional convolution operator, with optional dilation.
 
   Wraps XLA's `Conv
@@ -174,7 +168,7 @@ def conv_general_dilated(
 
 def conv(lhs: Array, rhs: Array, window_strides: Sequence[int],
          padding: str, precision: lax.PrecisionLike = None,
-         preferred_element_type: Optional[DType] = None) -> Array:
+         preferred_element_type: Optional[DTypeLike] = None) -> Array:
   """Convenience wrapper around `conv_general_dilated`.
 
   Args:
@@ -204,7 +198,7 @@ def conv_with_general_padding(lhs: Array, rhs: Array,
                               lhs_dilation: Optional[Sequence[int]],
                               rhs_dilation: Optional[Sequence[int]],
                               precision: lax.PrecisionLike = None,
-                              preferred_element_type: Optional[DType] = None) -> Array:
+                              preferred_element_type: Optional[DTypeLike] = None) -> Array:
   """Convenience wrapper around `conv_general_dilated`.
 
   Args:
@@ -256,7 +250,7 @@ def _conv_transpose_padding(k, s, padding):
     else:
       pad_a = int(np.ceil(pad_len / 2))
   elif padding == 'VALID':
-    pad_len = k + s - 2 + _max(k - s, 0)
+    pad_len = k + s - 2 + max(k - s, 0)
     pad_a = k - 1
   else:
     raise ValueError('Padding mode must be `SAME` or `VALID`.')
@@ -277,7 +271,7 @@ def conv_transpose(lhs: Array, rhs: Array, strides: Sequence[int],
                    dimension_numbers: ConvGeneralDilatedDimensionNumbers = None,
                    transpose_kernel: bool = False,
                    precision: lax.PrecisionLike = None,
-                   preferred_element_type: Optional[DType] = None) -> Array:
+                   preferred_element_type: Optional[DTypeLike] = None) -> Array:
   """Convenience wrapper for calculating the N-d convolution "transpose".
 
   This function directly calculates a fractionally strided conv rather than
@@ -343,7 +337,7 @@ def conv_transpose(lhs: Array, rhs: Array, strides: Sequence[int],
   if transpose_kernel:
     # flip spatial dims and swap input / output channel axes
     rhs = _flip_axes(rhs, np.array(dn.rhs_spec)[2:])
-    rhs = np.swapaxes(rhs, dn.rhs_spec[0], dn.rhs_spec[1])
+    rhs = np.swapaxes(rhs, dn.rhs_spec[0], dn.rhs_spec[1])  # type: ignore[assignment]
   return conv_general_dilated(lhs, rhs, one, pads, strides, rhs_dilation, dn,
                               precision=precision,
                               preferred_element_type=preferred_element_type)

jax/_src/lax/windowed_reductions.py

@@ -36,12 +36,11 @@
 from jax._src.lib.mlir import ir
 from jax._src.lib.mlir.dialects import hlo
 from jax._src.numpy.ufuncs import logaddexp
+from jax._src.typing import Array
 
 map = util.safe_map
 zip = util.safe_zip
 
-Array = Any
-
 
 def reduce_window(operand, init_value, computation: Callable,
                   window_dimensions: core.Shape, window_strides: Sequence[int],

jax/_src/nn/functions.py

@@ -28,9 +28,9 @@
 from jax._src import dtypes
 from jax._src import util
 from jax._src.core import AxisName
+from jax._src.typing import Array, ArrayLike
 from jax._src.ops.special import logsumexp as _logsumexp
 
-Array = Any
 
 # activations
 
@@ -165,7 +165,7 @@ def log_sigmoid(x: Array) -> Array:
   return -softplus(-x)
 
 @jax.jit
-def elu(x: Array, alpha: Array = 1.0) -> Array:
+def elu(x: Array, alpha: ArrayLike = 1.0) -> Array:
   r"""Exponential linear unit activation function.
 
   Computes the element-wise function:
@@ -190,7 +190,7 @@ def elu(x: Array, alpha: Array = 1.0) -> Array:
   return jnp.where(x > 0, x, alpha * jnp.expm1(safe_x))
 
 @jax.jit
-def leaky_relu(x: Array, negative_slope: Array = 1e-2) -> Array:
+def leaky_relu(x: Array, negative_slope: ArrayLike = 1e-2) -> Array:
   r"""Leaky rectified linear unit activation function.
 
   Computes the element-wise function:
@@ -237,7 +237,7 @@ def hard_tanh(x: Array) -> Array:
   return jnp.where(x > 1, 1, jnp.where(x < -1, -1, x))
 
 @jax.jit
-def celu(x: Array, alpha: Array = 1.0) -> Array:
+def celu(x: Array, alpha: ArrayLike = 1.0) -> Array:
   r"""Continuously-differentiable exponential linear unit activation.
 
   Computes the element-wise function:
@@ -431,14 +431,14 @@ def softmax(x: Array,
     :func:`log_softmax`
   """
   if jax.config.jax_softmax_custom_jvp:
-    return _softmax(x, axis, where, initial)
+    return _softmax(x, axis, where, initial)  # type: ignore[return-value]
   else:
     return _softmax_deprecated(x, axis, where, initial)
 
 # TODO(mattjj): replace softmax with _softmax when deprecation flag is removed
 @partial(jax.custom_jvp, nondiff_argnums=(1,))
 def _softmax(
-    x,
+    x: Array,
     axis: Optional[Union[int, tuple[int, ...]]] = -1,
     where: Optional[Array] = None,
     initial: Optional[Array] = None) -> Array:
@@ -455,7 +455,11 @@ def _softmax_jvp(axis, primals, tangents):
   y = _softmax(x, axis, where, initial)
   return y, y * (x_dot - (y * x_dot).sum(axis, where=where, keepdims=True))
 
-def _softmax_deprecated(x, axis, where, initial):
+def _softmax_deprecated(
+    x: Array,
+    axis: Optional[Union[int, tuple[int, ...]]] = -1,
+    where: Optional[Array] = None,
+    initial: Optional[Array] = None) -> Array:
   x_max = jnp.max(x, axis, where=where, initial=initial, keepdims=True)
   unnormalized = jnp.exp(x - lax.stop_gradient(x_max))
   result = unnormalized / jnp.sum(unnormalized, axis, where=where, keepdims=True)
@@ -469,7 +473,7 @@ def standardize(x: Array,
               axis: Optional[Union[int, tuple[int, ...]]] = -1,
               mean: Optional[Array] = None,
               variance: Optional[Array] = None,
-              epsilon: Array = 1e-5,
+              epsilon: ArrayLike = 1e-5,
               where: Optional[Array] = None) -> Array:
   r"""Normalizes an array by subtracting ``mean`` and dividing by :math:`\sqrt{\mathrm{variance}}`."""
   if mean is None:
@@ -484,11 +488,11 @@ def standardize(x: Array,
   return (x - mean) * lax.rsqrt(variance + epsilon)
 
 def normalize(x: Array,
-              axis: Optional[Union[int, tuple[int, ...]]] = -1,
-              mean: Optional[Array] = None,
-              variance: Optional[Array] = None,
-              epsilon: Array = 1e-5,
-              where: Optional[Array] = None) -> Array:
+            axis: Optional[Union[int, tuple[int, ...]]] = -1,
+            mean: Optional[Array] = None,
+            variance: Optional[Array] = None,
+            epsilon: ArrayLike = 1e-5,
+            where: Optional[Array] = None) -> Array:
   r"""Normalizes an array by subtracting ``mean`` and dividing by :math:`\sqrt{\mathrm{variance}}`."""
   warnings.warn("jax.nn.normalize will be deprecated. Use jax.nn.standardize instead.", DeprecationWarning)
   return standardize(x, axis, mean, variance, epsilon, where)

jax/_src/nn/initializers.py

@@ -23,18 +23,17 @@
 
 import numpy as np
 
-import jax
 import jax.numpy as jnp
 from jax import lax
 from jax import random
 from jax._src import core
 from jax._src import dtypes
+from jax._src.typing import Array, ArrayLike
 from jax._src.util import set_module
 
 export = set_module('jax.nn.initializers')
 
-KeyArray = jax.Array
-Array = Any
+KeyArray = Array
 # TODO: Import or define these to match
 # https://github.com/numpy/numpy/blob/main/numpy/typing/_dtype_like.py.
 DTypeLikeFloat = Any
@@ -240,7 +239,7 @@ def _complex_uniform(key: KeyArray,
   theta = 2 * jnp.pi * random.uniform(key_theta, shape, real_dtype).astype(dtype)
   return r * jnp.exp(1j * theta)
 
-def _complex_truncated_normal(key: KeyArray, upper: Array,
+def _complex_truncated_normal(key: KeyArray, upper: ArrayLike,
                               shape: Union[Sequence[int], core.NamedShape],
                               dtype: DTypeLikeInexact) -> Array:
   """

jax/_src/ops/scatter.py

@@ -16,7 +16,7 @@
 
 from collections.abc import Sequence
 import sys
-from typing import Any, Callable, Optional, Union
+from typing import Callable, Optional, Union
 import warnings
 
 import numpy as np
@@ -31,9 +31,9 @@
 from jax._src.numpy import lax_numpy as jnp
 from jax._src.numpy import reductions
 from jax._src.numpy.util import check_arraylike, promote_dtypes
+from jax._src.typing import Array
 
 
-Array = Any
 if sys.version_info >= (3, 10):
     from types import EllipsisType
     SingleIndex = Union[None, int, slice, Sequence[int], Array, EllipsisType]

jax/_src/scipy/optimize/_lbfgs.py

@@ -12,18 +12,19 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """The Limited-Memory Broyden-Fletcher-Goldfarb-Shanno minimization algorithm."""
-from typing import Any, Callable, NamedTuple, Optional, Union
+from typing import Callable, NamedTuple, Optional, Union
 from functools import partial
 
 import jax
 import jax.numpy as jnp
 from jax import lax
 from jax._src.scipy.optimize.line_search import line_search
+from jax._src.typing import Array
+
 
 _dot = partial(jnp.dot, precision=lax.Precision.HIGHEST)
 
 
-Array = Any
 
 class LBFGSResults(NamedTuple):
   """Results from L-BFGS optimization

jax/_src/state/types.py

@@ -23,14 +23,13 @@
 from jax._src import effects
 from jax._src import pretty_printer as pp
 from jax._src.util import safe_map, safe_zip
+from jax._src.typing import Array
 
 ## JAX utilities
 
 map, unsafe_map = safe_map, map
 zip, unsafe_zip = safe_zip, zip
 
-Array = Any
-
 _ref_effect_color = pp.Color.GREEN
 
 class RefEffect(effects.JaxprInputEffect):