{"record":{"id":"1defceae284a07c6","repo":"jax-ml/jax","slug":"only-take-along-axis-like-gathers-supported","errorCode":null,"errorMessage":"Only take_along_axis-like gathers supported","messagePattern":"Only take_along_axis-like gathers supported","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/lowering.py","lineNumber":3366,"sourceCode":"    indices,\n    *,\n    dimension_numbers,\n    slice_sizes,\n    unique_indices,\n    indices_are_sorted,\n    mode,\n    fill_value,\n):\n  in_aval = ctx.avals_in[0]\n  indices_aval = ctx.avals_in[1]\n  out_aval = ctx.avals_out[0]\n\n  if (\n      len(in_aval.shape) != len(out_aval.shape)\n      or indices_aval.shape[:-1] != out_aval.shape\n      or indices_aval.shape[-1] != 1\n  ):\n    raise NotImplementedError(\"Only take_along_axis-like gathers supported\")\n  rank = len(out_aval.shape)\n\n  # During lowering jnp.take_along_axis to lax.gather, we append extra dimension\n  # to the end of the indices array. We should reshape it back to the original\n  # shape before lowering to Mosaic and rely on MLIR canonicalization to remove\n  # the reshapes.\n  recovered_indices = vector.shape_cast(\n      ir.VectorType.get(\n          ctx.lowering_context.dynamic_shape_replacement_fn(out_aval.shape),\n          indices.type.element_type,\n      ),\n      indices,\n  )\n  # Note: current support for lax.gather is still very limited.\n  del fill_value\n\n  (\n      offset_dims,","sourceCodeStart":3348,"sourceCodeEnd":3384,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/lowering.py#L3348-L3384","documentation":"Mosaic (TPU Pallas) lowering rule for lax.gather only supports gathers shaped exactly like jnp.take_along_axis: the indices array must have the same leading shape as the output with a single trailing index dimension, and input/output ranks must match. Any other gather configuration (e.g. lax.gather with fancy multi-dimensional index maps) is rejected at compile time because the TPU lowering only implements the take_along_axis pattern.","triggerScenarios":"Calling jnp.take_along_axis or lax.gather inside a jax.pallas kernel with indices whose shape is not out_shape + (1,), mismatched input/output ranks, or a raw lax.gather with a custom index_dims/feature_bits configuration.","commonSituations":"Porting NumPy fancy-indexing (x[idx] with multi-dim idx) into a Pallas TPU kernel; reshaping indices incorrectly before take_along_axis; using lax.gather_with_default_batch_dims or gather with slice_sizes that don't reduce to the take_along_axis form.","solutions":["Reshape the operation to jnp.take_along_axis(x, indices[..., None], axis=axis) so indices shape equals out_shape + (1,)","Move the fancy gather out of the Pallas kernel and perform it on the host/JAX side before passing data in","Rewrite the gather as explicit indexing in a loop or use tpu.dynamic_gather-compatible single-axis gather","Check for a newer JAX version where gather support in Mosaic was extended"],"exampleFix":"// before\nout = lax.gather(x, idx_map, dimension_numbers=..., slice_sizes=...)\n// after\nout = jnp.take_along_axis(x, idx[..., None], axis=-1)","handlingStrategy":"validation","validationCode":"import jax.numpy as jnp\ndef check_take_along_shape(x_shape, idx_shape, axis):\n    out_shape = list(x_shape); out_shape.pop(axis)\n    return len(idx_shape) == len(out_shape) + 1 and idx_shape[-1] == 1 and idx_shape[:-1] == tuple(out_shape)","typeGuard":null,"tryCatchPattern":"catch NotImplementedError around kernel compilation and fall back to running the gather in plain JAX outside the kernel","preventionTips":["Use jnp.take_along_axis with idx[..., None] as the canonical gather form in kernels","Avoid lax.gather with custom dimension_numbers in Pallas code","Test kernels with jax.eval_shape before compiling"],"tags":["jax","pallas","tpu","gather","not-implemented"],"backgroundTag":"unsupported-operation-lowering","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}