{"record":{"id":"c724e9b839d5ea95","repo":"jax-ml/jax","slug":"one-hot-indexing-is-only-supported-for-up-to-50-le","errorCode":null,"errorMessage":"One Hot indexing is only supported for up to 50 leading dimensions.","messagePattern":"One Hot indexing is only supported for up to 50 leading dimensions\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/numpy/indexing.py","lineNumber":903,"sourceCode":"  if axis_size == 0:\n    return lax.full(out_shape, 0, a.dtype)\n\n  index_dtype = lax_utils.index_dtype_for_axis_size(\n      dtypes.dtype(indices), axis_size, wrap_negative_indices\n  )\n  indices = lax.convert_element_type(indices, index_dtype)\n\n  if wrap_negative_indices:\n    indices = _normalize_index(indices, axis_size)\n\n  if mode == \"one_hot\":\n    from jax import nn  # pyrefly: ignore[missing-import]\n\n    hot = nn.one_hot(indices, axis_size, dtype=np.bool_)\n    if a.ndim == 1:\n      return einsum.einsum(\"...b,b->...\", hot, a, preferred_element_type=a.dtype)\n    if axis_int > len(string.ascii_letters) - 2:\n      raise ValueError(\n          \"One Hot indexing is only supported for up to 50 leading dimensions.\"\n      )\n    labels = \"\".join([string.ascii_letters[i] for i in range(axis_int)])\n    eq = labels + \"y...z,\" + labels + \"z...->\" + labels + \"y...\"\n    return einsum.einsum(\n        eq,\n        hot,\n        a,\n        precision=lax.Precision.HIGHEST,\n        preferred_element_type=a.dtype,\n    )\n\n  index_dims = [i for i, idx in enumerate(idx_shape) if i == axis_int or not core.definitely_equal(idx, 1)]\n\n  gather_index_shape = tuple(np.array(out_shape)[index_dims]) + (1,)\n  gather_indices = lax.reshape(indices, gather_index_shape)\n  slice_sizes = []\n  offset_dims = []","sourceCodeStart":885,"sourceCodeEnd":921,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/numpy/indexing.py#L885-L921","documentation":"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.","triggerScenarios":"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.","commonSituations":"Extremely high-rank tensors from excessive stacking/vmap nesting; pathological shapes from broadcasting bugs that inflated rank.","solutions":["Reduce array rank: move/reshape axes so the take axis is early or the array is processed per-slice","Take with a reshaped 2-D view: flatten all non-taken axes into one, take, then reshape back","Fix upstream logic creating absurdly high-rank arrays (often a vmap/stack bug)"],"exampleFix":"// before\ny = jnp.take_along_axis(a, idx, axis=60)  # rank-61 array\n// after\nlead = int(np.prod(a.shape[:60]))\ny = jnp.take_along_axis(a.reshape(lead, -1), idx.reshape(lead, -1), axis=1).reshape(a.shape)","handlingStrategy":"validation","validationCode":"axis_int = axis if axis >= 0 else a.ndim + axis\nassert axis_int <= 50 or a.ndim == 1, 'take_along_axis supports at most 50 leading dims'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Keep take axes early in high-rank arrays","Flatten leading axes before take_along_axis for huge ranks","Investigate rank growth from vmap/stack bugs"],"tags":["jax","take-along-axis","rank-limit","einsum"],"backgroundTag":"array-rank-limit","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}