{"record":{"id":"a1a01fa9abc1a481","repo":"jax-ml/jax","slug":"take-along-axis-indices-must-be-1d-if-axis-none-g","errorCode":null,"errorMessage":"take_along_axis indices must be 1D if axis=None, got shape {}","messagePattern":"take_along_axis indices must be 1D if axis=None, got shape (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/numpy/indexing.py","lineNumber":868,"sourceCode":"\n    >>> idx = jnp.argmin(x, axis=1, keepdims=True)\n    >>> idx\n    Array([[1],\n           [0]], dtype=int32)\n    >>> jnp.take_along_axis(x, idx, axis=1)\n    Array([[3],\n           [2]], dtype=int32)\n  \"\"\"\n  a, indices = util.ensure_arraylike(\"take_along_axis\", arr, indices)\n  index_dtype = indices.dtype\n  idx_shape = np.shape(indices)\n  if not dtypes.issubdtype(index_dtype, np.integer):\n    raise TypeError(\"take_along_axis indices must be of integer type, got \"\n                    f\"{index_dtype}\")\n  if axis is None:\n    if np.ndim(indices) != 1:\n      msg = \"take_along_axis indices must be 1D if axis=None, got shape {}\"\n      raise ValueError(msg.format(idx_shape))\n    a = a.ravel()\n    axis = 0\n  rank = a.ndim\n  if rank != np.ndim(indices):\n    msg = \"indices and arr must have the same number of dimensions; {} vs. {}\"\n    raise ValueError(msg.format(np.ndim(indices), a.ndim))\n  axis_int = canonicalize_axis(axis, rank)\n\n  def replace(tup, val):\n    lst = list(tup)\n    lst[axis_int] = val\n    return tuple(lst)\n\n  axis_size = a.shape[axis_int]\n  arr_shape = replace(a.shape, 1)\n  out_shape = lax.broadcast_shapes(idx_shape, arr_shape)\n  if axis_size == 0:\n    return lax.full(out_shape, 0, a.dtype)","sourceCodeStart":850,"sourceCodeEnd":886,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/numpy/indexing.py#L850-L886","documentation":"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.","triggerScenarios":"jnp.take_along_axis(a, indices, axis=None) with np.ndim(indices) != 1 (e.g. shape (n, 1) or (n, m)).","commonSituations":"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.","solutions":["Flatten indices: indices = indices.ravel()","Specify an explicit axis instead of None when indices are multi-dimensional"],"exampleFix":"// before\ny = jnp.take_along_axis(a, idx_2d, axis=None)\n// after\ny = jnp.take_along_axis(a, idx_2d.ravel(), axis=None)","handlingStrategy":"validation","validationCode":"if axis is None:\n    assert indices.ndim == 1, f'indices must be 1D for axis=None, got {indices.shape}'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["ravel indices when using axis=None","Pass explicit axis for multi-dimensional index grids"],"tags":["jax","take-along-axis","shape-mismatch","axis"],"backgroundTag":"index-shape-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}