jax-ml/jax · error · ValueError

take_along_axis indices must be 1D if axis=None, got shape {

Error message

take_along_axis indices must be 1D if axis=None, got shape {}

What it means

When axis=None, take_along_axis flattens the input, so indices must be 1-D. Multi-dimensional indices are rejected because there's no well-defined mapping to a flattened gather.

Source

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

    >>> 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)

  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)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Flatten indices: indices = indices.ravel()
2. Specify an explicit axis instead of None when indices are multi-dimensional

Example fix

// before
y = jnp.take_along_axis(a, idx_2d, axis=None)
// after
y = jnp.take_along_axis(a, idx_2d.ravel(), axis=None)

if axis is None:
    assert indices.ndim == 1, f'indices must be 1D for axis=None, got {indices.shape}'

When it happens

Trigger: jnp.take_along_axis(a, indices, axis=None) with np.ndim(indices) != 1 (e.g. shape (n, 1) or (n, m)).

Common situations: Defaulting axis=None while passing 2-D index grids built for per-axis takes; forgetting to ravel indices when switching from axis= to axis=None.

Related errors

• indices and arr must have the same number of dimensions; {}
• 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
• Incorrect output shape for return value #{i}: Expected: {out

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