Fix type annotations for NumPy 2.2

Author: jakevdp

.pre-commit-config.yaml

@@ -36,7 +36,7 @@ repos:
   - id: mypy
     files: (jax/|tests/typing_test\.py)
     exclude: jax/_src/basearray.py|jax/numpy/__init__.py  # Use pyi instead
-    additional_dependencies: [types-requests==2.31.0, jaxlib, numpy~=2.1.0]
+    additional_dependencies: [types-requests==2.31.0, jaxlib, numpy>=2.2.0]
     args: [--config=pyproject.toml]
 
 - repo: https://github.com/mwouts/jupytext

jax/_src/dtypes.py

@@ -419,7 +419,7 @@ def _issubdtype_cached(a: type | np.dtype | ExtendedDType,
     return b_sctype in {a_sctype, np.unsignedinteger, np.integer, np.number, np.generic}
 
   # Otherwise, fall back to numpy.issubdtype
-  return np.issubdtype(a_sctype, b_sctype)
+  return bool(np.issubdtype(a_sctype, b_sctype))
 
 can_cast = np.can_cast
 

jax/_src/interpreters/mlir.py

@@ -2801,7 +2801,7 @@ def _wrapped_callback(*args):  # pylint: disable=function-redefined
 def _layout_to_mlir_layout(minor_to_major: Sequence[int] | None):
   if minor_to_major is None:
     # Needed for token layouts
-    layout = np.zeros((0,), dtype="int64")
+    layout: np.ndarray = np.zeros((0,), dtype="int64")
   else:
     layout = np.array(minor_to_major, dtype="int64")
   return ir.DenseIntElementsAttr.get(layout, type=ir.IndexType.get())

jax/_src/mesh_utils.py

@@ -386,7 +386,7 @@ def _create_device_mesh_for_nd_torus_splitting_axes(
           )
       ):
         best_logical_axis_assignment = logical_axis_assignment
-    assignment[:, logical_axis] = best_logical_axis_assignment
+    assignment[:, logical_axis] = best_logical_axis_assignment  # type: ignore  # numpy 2.2
 
   # Read out the assignment.
   logical_mesh = _generate_logical_mesh(
@@ -597,10 +597,10 @@ def _generate_logical_mesh(
           zip(logical_indices, physical_indices, range(len(logical_indices)))
       )
   )
-  logical_mesh = np.transpose(logical_mesh, transpose_axes)
+  logical_mesh = np.transpose(logical_mesh, transpose_axes)  # type: ignore  # numpy 2.2
 
   # Reshape to add the trivial dimensions back.
-  logical_mesh = np.reshape(logical_mesh, logical_mesh_shape)
+  logical_mesh = np.reshape(logical_mesh, logical_mesh_shape)  # type: ignore  # numpy 2.2
 
   return logical_mesh
 

jax/_src/numpy/linalg.py

@@ -1374,7 +1374,7 @@ def _lstsq(a: ArrayLike, b: ArrayLike, rcond: float | None, *,
     x = jnp.empty((n, *b.shape[1:]), dtype=a.dtype)
   else:
     if rcond is None:
-      rcond = jnp.finfo(dtype).eps * max(n, m)
+      rcond = float(jnp.finfo(dtype).eps) * max(n, m)
     else:
       rcond = jnp.where(rcond < 0, jnp.finfo(dtype).eps, rcond)
     u, s, vt = svd(a, full_matrices=False)

jax/_src/numpy/polynomial.py

@@ -246,7 +246,7 @@ def polyfit(x: ArrayLike, y: ArrayLike, deg: int, rcond: float | None = None,
 
   # set rcond
   if rcond is None:
-    rcond = len(x_arr) * finfo(x_arr.dtype).eps
+    rcond = len(x_arr) * float(finfo(x_arr.dtype).eps)
   rcond = core.concrete_or_error(float, rcond, "rcond must be float")
   # set up least squares equation for powers of x
   lhs = vander(x_arr, order)

jax/_src/op_shardings.py

@@ -100,7 +100,7 @@ def op_sharding_to_numpy_indices(
 
   for i, idxs in enumerate(itertools.product(*axis_indices)):
     for _ in range(num_replicas):
-      indices[next(device_it)] = idxs
+      indices[next(device_it)] = idxs  # type: ignore  # numpy 2.2
   return indices
 
 

jax/_src/sharding_impls.py

@@ -738,7 +738,7 @@ def __repr__(self) -> str:
     ids = self._ids.copy()
     platform_name = self._devices[0].platform.upper()
     for idx, x in np.ndenumerate(ids):
-      ids[idx] = DeviceIdSet(platform_name, *(self._devices[i].id for i in x))
+      ids[idx] = DeviceIdSet(platform_name, *(self._devices[i].id for i in x))  # type: ignore  # numpy 2.2
     body = np.array2string(ids, prefix=cls_name + '(', suffix=')',
                            max_line_width=100)
     mem = '' if self._memory_kind is None else f', memory_kind={self._memory_kind}'

jax/_src/sharding_specs.py

@@ -97,7 +97,7 @@ def _sharding_spec_indices(self, shape: tuple[int, ...]) -> np.ndarray:
   # is used to extract the corresponding shard of the logical array.
   shard_indices = np.empty([math.prod(shard_indices_shape)], dtype=np.object_)
   for i, idxs in enumerate(itertools.product(*axis_indices)):
-    shard_indices[i] = idxs
+    shard_indices[i] = idxs  # type: ignore  # numpy 2.2
   shard_indices = shard_indices.reshape(shard_indices_shape)
 
   # Ensure that each sharded axis is used exactly once in the mesh mapping

jax/experimental/sparse/bcoo.py

@@ -578,7 +578,7 @@ def _bcoo_transpose_transpose(ct, data, indices, *, permutation: Sequence[int],
     raise ValueError("Cannot transpose with respect to sparse indices")
   assert data_ct.dtype == data.aval.dtype
   ct_spinfo = SparseInfo(tuple(spinfo.shape[p] for p in permutation))
-  rev_permutation = list(np.argsort(permutation))
+  rev_permutation = list(map(int, np.argsort(permutation)))
   # TODO(jakevdp) avoid dummy indices?
   dummy_indices = jnp.zeros([1 for i in range(indices.ndim - 2)] + list(indices.shape[-2:]), dtype=int)
   data_trans, _ = _bcoo_transpose(data_ct, dummy_indices, permutation=rev_permutation, spinfo=ct_spinfo)
@@ -865,7 +865,7 @@ def _bcoo_dot_general_transpose(ct, lhs_data, lhs_indices, rhs, *, dimension_num
     dims: DotDimensionNumbers = ((ans_rhs, rhs_kept), (ans_batch, rhs_batch))
     lhs_contract_sorted_by_rhs = list(np.take(lhs_contract, np.argsort(rhs_contract)))
     permutation = list(lhs_batch) + lhs_kept + lhs_contract_sorted_by_rhs
-    out_axes = list(np.argsort(permutation))
+    out_axes = list(map(int, np.argsort(permutation)))
 
     # Determine whether efficient approach is possible:
     placeholder_data = jnp.empty((lhs_indices.ndim - 2) * (1,) + (lhs_indices.shape[-2],))