jax-ml/jax · error · ValueError

indices and arr must have the same number of dimensions; {}

Error message

indices and arr must have the same number of dimensions; {} vs. {}

What it means

take_along_axis requires indices and arr to have identical rank so each output element has a well-defined source. The message reports ndim(indices) vs ndim(arr).

Source

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

    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)

  axis_size = a.shape[axis_int]
  arr_shape = replace(a.shape, 1)
  out_shape = lax.broadcast_shapes(idx_shape, arr_shape)
  if axis_size == 0:
    return lax.full(out_shape, 0, a.dtype)

  index_dtype = lax_utils.index_dtype_for_axis_size(
      dtypes.dtype(indices), axis_size, wrap_negative_indices
  )
  indices = lax.convert_element_type(indices, index_dtype)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Add the missing axis: indices = indices[..., None] when indices lacks the taken dimension
2. Or reshape a to match: a = a.reshape(...) if the intent is a flattened take
3. Verify both ndims with an assert during development

Example fix

// before
# a: (B, N, D), idx: (B, N)
y = jnp.take_along_axis(a, idx, axis=2)
// after
y = jnp.take_along_axis(a, idx[..., None], axis=2)

assert a.ndim == indices.ndim, f'rank mismatch: a={a.ndim}, idx={indices.ndim}'

When it happens

Trigger: Calling jnp.take_along_axis(a, indices, axis=k) where a.ndim != indices.ndim — e.g. a is (B, N, D) and indices is (B, N) without a trailing dimension.

Common situations: Forgetting to add a trailing axis of size 1 to indices to broadcast along the taken dimension; mixing ranks after slicing operations that drop dims.

Related errors

• take_along_axis indices must be 1D if axis=None, got shape {
• put_along_axis arguments 'arr' and 'indices' must have same
• unexpected JAX type (e.g. shape/dtype) for argument to VJP f
• cotangent type does not match function output, expected {out
• Mismatched number of outputs from callback. Expected: {}, Ac

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