{"record":{"id":"d7221668c4729245","repo":"jax-ml/jax","slug":"take-along-axis-indices-must-be-of-integer-type-g","errorCode":null,"errorMessage":"take_along_axis indices must be of integer type, got {index_dtype}","messagePattern":"take_along_axis indices must be of integer type, got (.+?)","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/numpy/indexing.py","lineNumber":863,"sourceCode":"    Array([[3, 4, 5],\n           [2, 6, 7]], dtype=int32)\n\n    Similarly, we can use :func:`~jax.numpy.argmin` with ``keepdims=True`` and\n    use ``take_along_axis`` to extract the minimum value:\n\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","sourceCodeStart":845,"sourceCodeEnd":881,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/numpy/indexing.py#L845-L881","documentation":"jnp.take_along_axis requires integer (or boolean-as-integer unsupported) index arrays; float indices cannot be lowered to a gather. The actual dtype is reported in the message.","triggerScenarios":"Calling jnp.take_along_axis(a, indices, axis) where indices has a float dtype (e.g. results of argsort-free computations, float argsort outputs from other libs, or jnp.arange defaults in some dtypes contexts).","commonSituations":"Indices produced by float math or another library (NumPy argsort returns int, but e.g. some distance-argmin pipelines yield floats); forgetting .astype(int) after computing positions.","solutions":["Cast indices: indices.astype(jnp.int32) (or int64 per platform)","Fix the producing computation so it emits integer indices (e.g. use jnp.argsort / argmin outputs directly)"],"exampleFix":"// before\npositions = some_float_computation(...)\ny = jnp.take_along_axis(a, positions, axis=1)\n// after\ny = jnp.take_along_axis(a, positions.astype(jnp.int32), axis=1)","handlingStrategy":"type-guard","validationCode":"assert jnp.issubdtype(indices.dtype, jnp.integer), f'indices must be int, got {indices.dtype}'","typeGuard":"def is_integer_indices(idx) -> bool:\n    import jax.numpy as jnp\n    return jnp.issubdtype(idx.dtype, jnp.integer)","tryCatchPattern":"try:\n    y = jnp.take_along_axis(a, idx, axis=axis)\nexcept TypeError:\n    y = jnp.take_along_axis(a, idx.astype(jnp.int32), axis=axis)","preventionTips":["Always cast index arrays: idx.astype(jnp.int32)","Use integer-producing ops (argmax/argsort) for indices"],"tags":["jax","take-along-axis","dtype","integer-index"],"backgroundTag":"wrong-index-dtype","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}