jax-ml/jax · error · TypeError

take_along_axis indices must be of integer type, got {index_

Error message

take_along_axis indices must be of integer type, got {index_dtype}

What it means

jnp.take_along_axis requires integer (or boolean-as-integer unsupported) index arrays; float indices cannot be lowered to a gather. The actual dtype is reported in the message.

Source

Thrown at jax/_src/numpy/indexing.py:863

    Array([[3, 4, 5],
           [2, 6, 7]], dtype=int32)

    Similarly, we can use :func:`~jax.numpy.argmin` with ``keepdims=True`` and
    use ``take_along_axis`` to extract the minimum value:

    >>> idx = jnp.argmin(x, axis=1, keepdims=True)
    >>> idx
    Array([[1],
           [0]], dtype=int32)
    >>> jnp.take_along_axis(x, idx, axis=1)
    Array([[3],
           [2]], dtype=int32)
  """
  a, indices = util.ensure_arraylike("take_along_axis", arr, indices)
  index_dtype = indices.dtype
  idx_shape = np.shape(indices)
  if not dtypes.issubdtype(index_dtype, np.integer):
    raise TypeError("take_along_axis indices must be of integer type, got "
                    f"{index_dtype}")
  if axis is None:
    if np.ndim(indices) != 1:
      msg = "take_along_axis indices must be 1D if axis=None, got shape {}"
      raise ValueError(msg.format(idx_shape))
    a = a.ravel()
    axis = 0
  rank = a.ndim
  if rank != np.ndim(indices):
    msg = "indices and arr must have the same number of dimensions; {} vs. {}"
    raise ValueError(msg.format(np.ndim(indices), a.ndim))
  axis_int = canonicalize_axis(axis, rank)

  def replace(tup, val):
    lst = list(tup)
    lst[axis_int] = val
    return tuple(lst)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Cast indices: indices.astype(jnp.int32) (or int64 per platform)
2. Fix the producing computation so it emits integer indices (e.g. use jnp.argsort / argmin outputs directly)

Example fix

// before
positions = some_float_computation(...)
y = jnp.take_along_axis(a, positions, axis=1)
// after
y = jnp.take_along_axis(a, positions.astype(jnp.int32), axis=1)

assert jnp.issubdtype(indices.dtype, jnp.integer), f'indices must be int, got {indices.dtype}'

def is_integer_indices(idx) -> bool:
    import jax.numpy as jnp
    return jnp.issubdtype(idx.dtype, jnp.integer)

try:
    y = jnp.take_along_axis(a, idx, axis=axis)
except TypeError:
    y = jnp.take_along_axis(a, idx.astype(jnp.int32), axis=axis)

When it happens

Trigger: Calling jnp.take_along_axis(a, indices, axis) where indices has a float dtype (e.g. results of argsort-free computations, float argsort outputs from other libs, or jnp.arange defaults in some dtypes contexts).

Common situations: Indices produced by float math or another library (NumPy argsort returns int, but e.g. some distance-argmin pipelines yield floats); forgetting .astype(int) after computing positions.

Related errors

• primal and tangent arguments to jax.jvp do not match; dtypes
• unexpected JAX type (e.g. shape/dtype) for gradient ref pass
• Accumulator aval mismatch: expected {aval}, got {acc.aval}
• unexpected JAX type (e.g. shape/dtype) for argument to VJP f
• linear_transpose only supports [float or complex] -> [float

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/d7221668c4729245. Report an issue: GitHub.