jax-ml/jax · error · IndexError

Cannot do a non-empty jnp.take() from an empty axis.

Error message

Cannot do a non-empty jnp.take() from an empty axis.

What it means

You cannot gather elements from an axis of size 0 with a non-empty index array — there is nothing to take. JAX raises IndexError mirroring NumPy's empty-axis take failure.

Source

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

    gather_mode = slicing.GatherScatterMode.FILL_OR_DROP
    # lax.gather() does not support negative indices, so we wrap them here
    indices = util._where(indices < 0, indices + a.shape[axis_idx], indices)
  elif mode == "raise":
    # TODO(phawkins): we have no way to report out of bounds errors yet.
    raise NotImplementedError("The 'raise' mode to jnp.take is not supported.")
  elif mode == "wrap":
    indices = ufuncs.mod(indices, lax._const(indices, a.shape[axis_idx]))
    gather_mode = slicing.GatherScatterMode.PROMISE_IN_BOUNDS
  elif mode == "clip":
    gather_mode = slicing.GatherScatterMode.CLIP
  else:
    raise ValueError(f"Invalid mode '{mode}' for np.take")

  index_dims = len(np.shape(indices))
  slice_sizes = list(np.shape(a))
  if slice_sizes[axis_idx] == 0:
    if indices.size != 0:
      raise IndexError("Cannot do a non-empty jnp.take() from an empty axis.")
    return a

  if indices.size == 0:
    out_shape = (slice_sizes[:axis_idx] + list(indices.shape) +
                 slice_sizes[axis_idx + 1:])
    return lax.full_like(a, 0, shape=out_shape)

  slice_sizes[axis_idx] = 1
  dnums = slicing.GatherDimensionNumbers(
    offset_dims=tuple(
      list(range(axis_idx)) +
      list(range(axis_idx + index_dims, len(a.shape) + index_dims - 1))),
    collapsed_slice_dims=(axis_idx,),
    start_index_map=(axis_idx,))
  return slicing.gather(a, indices[..., None], dimension_numbers=dnums,
                        slice_sizes=tuple(slice_sizes),
                        mode=gather_mode, unique_indices=unique_indices,
                        indices_are_sorted=indices_are_sorted, fill_value=fill_value)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Guard: if a.shape[axis] == 0 or indices.size == 0, skip or return an empty result
2. Fix upstream emptiness (ensure the source array is non-empty before taking)
3. Use mode='fill' semantics won't help — fix the empty-axis input

Example fix

// before
y = jnp.take(a, idx)  # a may be empty
// after
y = jnp.take(a, idx) if a.size else jnp.array([], dtype=a.dtype)

if a.shape[axis] == 0:
    assert indices.size == 0, 'cannot take from empty axis with non-empty indices'

try:
    y = jnp.take(a, idx, axis=axis)
except IndexError:
    y = jnp.empty((0,), dtype=a.dtype)  # empty fallback

When it happens

Trigger: jnp.take(a, indices, axis=k) where a.shape[k] == 0 and indices.size > 0 (e.g. taking from an empty list/array).

Common situations: Empty datasets, empty vocabulary lookups, filtered collections that became empty; take over an axis that a previous operation collapsed to zero length.

Related errors

• index is out of bounds for axis {axis} with size 0
• reduction axes {axes} contains out-of-bounds indices for {op
• argmin and argmax require non-empty reduced dimension. opera
• axis argument out of range: {axis=} for {operand.shape=}
• index {i} out of bounds for axis {axis} with size {size} ({n

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