jax-ml/jax · error · ValueError

One Hot indexing is only supported for up to 50 leading dime

Error message

One Hot indexing is only supported for up to 50 leading dimensions.

What it means

When the array is multi-dimensional, take_along_axis uses a one-hot einsum trick that labels leading axes with ASCII letters; with 52 letters available, only up to 50 leading dimensions before the taken axis are supported.

Source

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

  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)

  if wrap_negative_indices:
    indices = _normalize_index(indices, axis_size)

  if mode == "one_hot":
    from jax import nn  # pyrefly: ignore[missing-import]

    hot = nn.one_hot(indices, axis_size, dtype=np.bool_)
    if a.ndim == 1:
      return einsum.einsum("...b,b->...", hot, a, preferred_element_type=a.dtype)
    if axis_int > len(string.ascii_letters) - 2:
      raise ValueError(
          "One Hot indexing is only supported for up to 50 leading dimensions."
      )
    labels = "".join([string.ascii_letters[i] for i in range(axis_int)])
    eq = labels + "y...z," + labels + "z...->" + labels + "y..."
    return einsum.einsum(
        eq,
        hot,
        a,
        precision=lax.Precision.HIGHEST,
        preferred_element_type=a.dtype,
    )

  index_dims = [i for i, idx in enumerate(idx_shape) if i == axis_int or not core.definitely_equal(idx, 1)]

  gather_index_shape = tuple(np.array(out_shape)[index_dims]) + (1,)
  gather_indices = lax.reshape(indices, gather_index_shape)
  slice_sizes = []
  offset_dims = []

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Reduce array rank: move/reshape axes so the take axis is early or the array is processed per-slice
2. Take with a reshaped 2-D view: flatten all non-taken axes into one, take, then reshape back
3. Fix upstream logic creating absurdly high-rank arrays (often a vmap/stack bug)

Example fix

// before
y = jnp.take_along_axis(a, idx, axis=60)  # rank-61 array
// after
lead = int(np.prod(a.shape[:60]))
y = jnp.take_along_axis(a.reshape(lead, -1), idx.reshape(lead, -1), axis=1).reshape(a.shape)

axis_int = axis if axis >= 0 else a.ndim + axis
assert axis_int <= 50 or a.ndim == 1, 'take_along_axis supports at most 50 leading dims'

When it happens

Trigger: Calling jnp.take_along_axis on arrays with rank > ~52 where the take axis is late (axis_int > 50), triggering the einsum-based path instead of the simple 1-D path.

Common situations: Extremely high-rank tensors from excessive stacking/vmap nesting; pathological shapes from broadcasting bugs that inflated rank.

Related errors

• Run out of characters for batch dimension in einsum.
• Transpose of Einsum with multiple linear inputs is not suppo
• take_along_axis indices must be of integer type, got {index_
• take_along_axis indices must be 1D if axis=None, got shape {
• indices and arr must have the same number of dimensions; {}

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