{"record":{"id":"1e368eabd5d7100e","repo":"jax-ml/jax","slug":"indices-and-arr-must-have-the-same-number-of-dimen","errorCode":null,"errorMessage":"indices and arr must have the same number of dimensions; {} vs. {}","messagePattern":"indices and arr must have the same number of dimensions; (.+?) vs\\. (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/numpy/indexing.py","lineNumber":874,"sourceCode":"    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)\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","sourceCodeStart":856,"sourceCodeEnd":892,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/numpy/indexing.py#L856-L892","documentation":"take_along_axis requires indices and arr to have identical rank so each output element has a well-defined source. The message reports ndim(indices) vs ndim(arr).","triggerScenarios":"Calling jnp.take_along_axis(a, indices, axis=k) where a.ndim != indices.ndim — e.g. a is (B, N, D) and indices is (B, N) without a trailing dimension.","commonSituations":"Forgetting to add a trailing axis of size 1 to indices to broadcast along the taken dimension; mixing ranks after slicing operations that drop dims.","solutions":["Add the missing axis: indices = indices[..., None] when indices lacks the taken dimension","Or reshape a to match: a = a.reshape(...) if the intent is a flattened take","Verify both ndims with an assert during development"],"exampleFix":"// before\n# a: (B, N, D), idx: (B, N)\ny = jnp.take_along_axis(a, idx, axis=2)\n// after\ny = jnp.take_along_axis(a, idx[..., None], axis=2)","handlingStrategy":"validation","validationCode":"assert a.ndim == indices.ndim, f'rank mismatch: a={a.ndim}, idx={indices.ndim}'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Add trailing axis to indices: idx[..., None]","Assert matching ranks in tests for index pipelines"],"tags":["jax","take-along-axis","ndim","shape-mismatch"],"backgroundTag":"index-shape-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}