Add common helpers index_dtype_for_shape and index_dtype_for_dim.

Use them consistently when choosing a dtype for an array index.

PiperOrigin-RevId: 800579624

Author: hawkinsp

jax/_src/lax/lax.py

@@ -57,6 +57,7 @@
 from jax._src.interpreters import pxla
 from jax._src.interpreters.batching import RaggedAxis
 from jax._src.lax import slicing
+from jax._src.lax import utils as lax_utils
 from jax._src.mesh import get_abstract_mesh, get_concrete_mesh
 from jax._src.lax.utils import (
   input_dtype, dtype_to_string, standard_abstract_eval,
@@ -7100,12 +7101,13 @@ def _squeeze_lower(ctx, operand, *, dimensions):
 
 def shape_as_value(shape: core.Shape):
   """Converts a shape that may contain Poly values into a JAX value."""
+  dtype = lax_utils.int_dtype_for_shape(shape, signed=True)
   if len(shape) == 0:
-    return full((0,), np.array(0, np.int64))
+    return full((0,), np.array(0, dtype=dtype))
   if core.is_constant_shape(shape):
-    return np.asarray(shape, dtype=np.int64)
+    return np.asarray(shape, dtype=dtype)
   dims = [
-      expand_dims(convert_element_type(core.dimension_as_value(d), np.int64),
+      expand_dims(convert_element_type(core.dimension_as_value(d), dtype),
                   (0,))
       for d in shape
   ]
@@ -8113,9 +8115,10 @@ def _operands_to_keys(*operands, num_keys=1):
 
 def _sort_jvp(primals, tangents, *, dimension, is_stable, num_keys):
   shape = primals[0].shape
+  index_dtype = lax_utils.int_dtype_for_shape(shape, signed=False)
   sorted_primals_and_idx = sort_p.bind(
       *primals,
-      broadcasted_iota(dtypes.canonicalize_dtype(np.uint64), shape, dimension),
+      broadcasted_iota(index_dtype, shape, dimension),
       dimension=dimension, is_stable=is_stable, num_keys=num_keys)
   batch_dims = tuple(np.delete(np.arange(len(shape), dtype=np.int64),
                                dimension))

jax/_src/lax/slicing.py

@@ -38,6 +38,7 @@
 from jax._src.interpreters import mlir
 from jax._src.interpreters import partial_eval as pe
 from jax._src.lax import lax
+from jax._src.lax import utils as lax_utils
 from jax._src.lax.utils import (
     _argnum_weak_type,
     input_dtype,
@@ -2107,16 +2108,19 @@ def _gather_fill(operand, indices, *, dimension_numbers, slice_sizes,
                  output_shape):
   """Lowers a FILL_OR_DROP gather as a PROMISE_IN_BOUNDS gather with masking."""
   dnums = dimension_numbers
-  intarray = partial(np.array, dtype=np.int64)
-  operand_dims = lax.shape_as_value(operand.shape)
-  indices = lax.convert_element_type(indices, np.int64)
+  index_dtype = lax_utils.int_dtype_for_shape(operand.shape, signed=True)
+  intarray = partial(np.array, dtype=index_dtype)
+  operand_dims = lax.shape_as_value(operand.shape).astype(index_dtype)
+  indices = lax.convert_element_type(indices, index_dtype)
   num_batch_dims = len(indices.shape) - 1
 
-  upper_bound = (
-      operand_dims[intarray(dnums.start_index_map)] -
-      lax.shape_as_value(slice_sizes)[intarray(dnums.start_index_map)])
+  upper_bound = operand_dims[
+      intarray(dnums.start_index_map)
+  ] - lax.shape_as_value(slice_sizes)[intarray(dnums.start_index_map)].astype(
+      index_dtype
+  )
   mask = lax.bitwise_and(
-      lax.ge(indices, np.int64(0)),
+      lax.ge(indices, index_dtype.type(0)),
       lax.le(indices, lax.expand_dims(upper_bound, tuple(range(num_batch_dims)))))
   mask = lax.reduce_and(mask, [num_batch_dims])
 
@@ -2727,18 +2731,20 @@ def _clamp_scatter_indices(operand, indices, updates, *, dnums):
 
   upper_bounds: core.Shape = tuple(operand.shape[i] - slice_sizes[i]
                                    for i in dnums.scatter_dims_to_operand_dims)
+
   # Stack upper_bounds into a Array[n]
   upper_bound = lax.shape_as_value(upper_bounds)
   # This fix fails lax_test_no_jax_array
-  upper_bound = lax.min(upper_bound,
-                        lax.convert_element_type(np.uint64(np.iinfo(indices.dtype).max),
-                                                  np.int64))
-
+  upper_bound = lax.min(
+      upper_bound,
+      upper_bound.dtype.type(
+          min(np.iinfo(upper_bound.dtype).max, np.iinfo(indices.dtype).max)
+      ),
+  )
+  upper_bound = lax.convert_element_type(upper_bound, indices.dtype)
   upper_bound = lax.broadcast_in_dim(upper_bound, indices.shape,
                                      (len(indices.shape) - 1,))
-  return lax.clamp(np.int64(0), lax.convert_element_type(indices, np.int64),
-                   upper_bound)
-
+  return lax.clamp(indices.dtype.type(0), indices, upper_bound)
 
 def _scatter_addsub_jvp(
     prim, primals, tangents, *, update_jaxpr, update_consts, dimension_numbers,
@@ -3132,10 +3138,7 @@ def _scatter_jvp(primals, tangents, *, update_jaxpr, update_consts,
   for update_dim in dnums.update_window_dims:
     ids_shape[update_dim] = 1
   num_ids = math.prod(ids_shape)
-  if core.is_constant_dim(num_ids):
-    id_dtype = np.uint32 if (num_ids + 1) < np.iinfo(np.uint32).max else np.uint64
-  else:
-    id_dtype = dtypes.canonicalize_dtype(np.uint64)
+  id_dtype = lax_utils.int_dtype_for_dim(num_ids, signed=False)
   update_ids = lax.add(lax.reshape(lax.iota(id_dtype, num_ids), ids_shape),
                        lax._ones(updates, dtype=id_dtype))
 

jax/_src/lax/utils.py

@@ -18,17 +18,20 @@
 
 from functools import partial
 
+import numpy as np
+
 from jax._src import core
 from jax._src import dispatch
+from jax._src import dtypes
 from jax._src import mesh as mesh_lib
 from jax._src import state
 from jax._src.named_sharding import DuplicateSpecError, NamedSharding
 from jax._src.partition_spec import PartitionSpec as P
 from jax._src.util import safe_zip
+from jax._src.typing import DimSize, DType, Shape
 
 zip, unsafe_zip = safe_zip, zip
 
-import numpy as np
 
 def input_dtype(x, *_, **__):
   return x.dtype
@@ -223,3 +226,36 @@ def dtype_to_string(dtype):
   except AttributeError:
     pass
   return str(dtype)
+
+_int32_max = np.iinfo(np.int32).max
+_uint32_max = np.iinfo(np.uint32).max
+
+def int_dtype_for_dim(d: DimSize, *, signed: bool) -> DType:
+  """Returns a integer dtype large enough to contain indices in dimension d."""
+  if signed:
+    if not core.is_constant_dim(d):
+      return dtypes.default_int_dtype()
+    return np.dtype(np.int64) if d > _int32_max else np.dtype(np.int32)
+  else:
+    if not core.is_constant_dim(d):
+      return dtypes.default_uint_dtype()
+    return np.dtype(np.uint64) if d > _uint32_max else np.dtype(np.uint32)
+
+def int_dtype_for_shape(shape: Shape, *, signed: bool) -> DType:
+  """Returns a integer dtype large enough to contain indices in `shape`."""
+  if signed:
+    for d in shape:
+      if core.is_constant_dim(d):
+        if d > _int32_max:
+          return np.dtype(np.int64)
+      else:
+        return dtypes.default_int_dtype()
+    return np.dtype(np.int32)
+  else:
+    for d in shape:
+      if core.is_constant_dim(d):
+        if d > _uint32_max:
+          return np.dtype(np.uint64)
+      else:
+        return dtypes.default_uint_dtype()
+    return np.dtype(np.uint32)

jax/_src/numpy/indexing.py

@@ -32,6 +32,7 @@
 from jax._src import errors
 from jax._src.lax import lax
 from jax._src.lax import slicing
+from jax._src.lax import utils as lax_utils
 from jax._src.numpy import einsum
 from jax._src.numpy import error as jnp_error
 from jax._src.numpy import lax_numpy
@@ -305,8 +306,7 @@ def replace(tup, val):
     lst[axis_int] = val
     return tuple(lst)
 
-  use_64bit_index = any(not core.is_constant_dim(d) or d >= (1 << 31) for d in a.shape)
-  index_dtype = np.dtype('int64' if use_64bit_index else 'int32')
+  index_dtype = lax_utils.int_dtype_for_dim(a.shape, signed=True)
   indices = lax.convert_element_type(indices, index_dtype)
 
   axis_size = a.shape[axis_int]
@@ -850,10 +850,7 @@ def index_to_gather(x_shape: Sequence[int], idx: Sequence[Any],
   collapsed_slice_dims: list[int] = []
   start_index_map: list[int] = []
 
-  use_64bit_index = (
-    any(not core.is_constant_dim(d) or d >= (1 << 31) for d in x_shape) and
-    config.enable_x64.value)
-  index_dtype = np.dtype('int64') if use_64bit_index else np.dtype('int32')
+  index_dtype = lax_utils.int_dtype_for_shape(x_shape, signed=True)
 
   # Gather indices.
   # Pairs of (array, start_dim) values. These will be broadcast into

jax/_src/numpy/lax_numpy.py

@@ -48,6 +48,7 @@
 from jax._src.lax import lax
 from jax._src.lax import slicing as lax_slicing
 from jax._src.lax import special as lax_special
+from jax._src.lax import utils as lax_utils
 from jax._src.lib import xla_client as xc
 from jax._src.numpy.array_constructors import array, asarray
 from jax._src.numpy import array_creation
@@ -7166,8 +7167,12 @@ def diag_indices(n: int, ndim: int = 2) -> tuple[Array, ...]:
   if ndim < 0:
     raise ValueError("ndim argument to diag_indices must be nonnegative, got {}"
                      .format(ndim))
-  # TODO(phawkins): Use an int64 index if n >= 2**31.
-  return (lax.iota(int, n),) * ndim
+  index_dtype = lax_utils.int_dtype_for_dim(n, signed=True)
+  # We'd give the correct output values with int32, but use the default dtype to
+  # match NumPy type semantics if x64 mode is enabled for now.
+  if index_dtype == np.dtype(np.int32):
+    index_dtype = dtypes.default_int_dtype()
+  return (lax.iota(index_dtype, n),) * ndim
 
 
 @export
@@ -9258,7 +9263,8 @@ def body_fun(state, _):
 
 
 def _searchsorted_via_sort(sorted_arr: Array, query: Array, side: str, dtype: type) -> Array:
-  working_dtype = np.dtype('int32') if sorted_arr.size + query.size < np.iinfo(np.int32).max else np.dtype('int64')
+  working_dtype = lax_utils.int_dtype_for_dim(sorted_arr.size + query.size,
+                                              signed=False)
   def _rank(x):
     idx = lax.iota(working_dtype, x.shape[0])
     return array_creation.zeros_like(idx).at[argsort(x)].set(idx)
@@ -9354,7 +9360,7 @@ def searchsorted(a: ArrayLike, v: ArrayLike, side: str = 'left',
   a, v = util.promote_dtypes(a, v)
   if sorter is not None:
     a = a[sorter]
-  dtype = np.dtype('int32') if a.shape[0] <= np.iinfo(np.int32).max else np.dtype('int64')
+  dtype = lax_utils.int_dtype_for_dim(a.shape[0], signed=True)
   if a.shape[0] == 0:
     return array_creation.zeros_like(v, dtype=dtype)
   impl = {

jax/_src/numpy/linalg.py

@@ -30,6 +30,7 @@
 from jax._src.custom_derivatives import custom_jvp
 from jax._src.lax import lax
 from jax._src.lax import linalg as lax_linalg
+from jax._src.lax import utils as lax_utils
 from jax._src.numpy import array_creation
 from jax._src.numpy import einsum
 from jax._src.numpy import indexing
@@ -294,11 +295,8 @@ def svd(
     s = lax.abs(v)
     if compute_uv:
       sign = lax.sign(v)
-      idx_dtype = (
-          np.int64
-          if int(s.shape[s.ndim - 1]) > np.iinfo(np.int32).max
-          else np.int32
-      )
+      idx_dtype = lax_utils.int_dtype_for_dim(
+          s.shape[s.ndim - 1], signed=False)
       idxs = lax.broadcasted_iota(idx_dtype, s.shape, dimension=s.ndim - 1)
       s, idxs, sign = lax.sort((s, idxs, sign), dimension=-1, num_keys=1)
       s = lax.rev(s, dimensions=[s.ndim - 1])

jax/_src/numpy/setops.py

@@ -26,6 +26,7 @@
 from jax._src import dtypes
 from jax._src.lax import lax
 from jax._src.lax import slicing as lax_slicing
+from jax._src.lax import utils as lax_utils
 from jax._src.numpy.array_creation import empty, full, full_like, ones, zeros
 from jax._src.numpy.lax_numpy import (
     append, arange, concatenate, diff,
@@ -344,11 +345,7 @@ def _intersect1d_sorted_mask(arr1: Array, arr2: Array,
   assert arr1.ndim == arr2.ndim == 1
   arr = concatenate((arr1, arr2))
   if return_indices:
-    use_64bit_index = (
-        not core.is_constant_dim(arr.shape[0])
-        or arr.shape[0] >= np.iinfo(np.int32).max
-    )
-    idx_dtype = np.int64 if use_64bit_index else np.int32
+    idx_dtype = lax_utils.int_dtype_for_dim(arr.shape[0], signed=True)
     iota = lax.broadcasted_iota(idx_dtype, np.shape(arr), dimension=0)
     aux, indices = lax.sort_key_val(arr, iota)
   else:

jax/_src/numpy/sorting.py

@@ -18,9 +18,9 @@
 import numpy as np
 
 from jax._src import api
-from jax._src import core
 from jax._src import dtypes
 from jax._src.lax import lax
+from jax._src.lax import utils as lax_utils
 from jax._src.numpy import util
 from jax._src.util import canonicalize_axis, set_module
 from jax._src.typing import Array, ArrayLike
@@ -154,11 +154,12 @@ def argsort(
     arr = arr.ravel()
     axis = 0
   dimension = canonicalize_axis(axis, arr.ndim)
-  use_64bit_index = core.is_constant_dim(arr.shape[dimension]) and arr.shape[dimension] >= (1 << 31)
-  # TODO(phawkins): we should probably use int64 indices for unknown dimensions,
-  # but that requires that we first support using int64 in a non-x64
-  # computation.
-  iota = lax.broadcasted_iota(np.dtype('int64') if use_64bit_index else int, arr.shape, dimension)
+  idx_dtype = lax_utils.int_dtype_for_dim(arr.shape[dimension], signed=True)
+  # We'd give the correct output values with int32, but use the default dtype to
+  # match NumPy type semantics if x64 mode is enabled for now.
+  if idx_dtype == np.dtype(np.int32):
+    idx_dtype = dtypes.default_int_dtype()
+  iota = lax.broadcasted_iota(idx_dtype, arr.shape, dimension)
   # For stable descending sort, we reverse the array and indices to ensure that
   # duplicates remain in their original order when the final indices are reversed.
   # For non-stable descending sort, we can avoid these extra operations.
@@ -425,7 +426,11 @@ def lexsort(keys: Array | np.ndarray | Sequence[ArrayLike], axis: int = -1) -> A
   if np.ndim(key_arrays[0]) == 0:
     return lax.full((), 0, dtypes.default_int_dtype())
   axis = canonicalize_axis(axis, np.ndim(key_arrays[0]))
-  use_64bit_index = key_arrays[0].shape[axis] >= (1 << 31)
-  iota = lax.broadcasted_iota(np.dtype('int64') if use_64bit_index else int,
-                              np.shape(key_arrays[0]), axis)
+  idx_dtype = lax_utils.int_dtype_for_dim(key_arrays[0].shape[axis],
+                                          signed=True)
+  # We'd give the correct output values with int32, but use the default dtype to
+  # match NumPy type semantics if x64 mode is enabled for now.
+  if idx_dtype == np.dtype(np.int32):
+    idx_dtype = dtypes.default_int_dtype()
+  iota = lax.broadcasted_iota(idx_dtype, np.shape(key_arrays[0]), axis)
   return lax.sort((*key_arrays[::-1], iota), dimension=axis, num_keys=len(key_arrays))[-1]