What causes "Only take_along_axis-like gathers supported" in jax-ml/jax?

Reshape the operation to jnp.take_along_axis(x, indices[..., None], axis=axis) so indices shape equals out_shape + (1,)

jax-ml/jax · error · NotImplementedError

Only take_along_axis-like gathers supported

Error message

Only take_along_axis-like gathers supported

What it means

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.

Source

Thrown at jax/_src/pallas/mosaic/lowering.py:3366

    indices,
    *,
    dimension_numbers,
    slice_sizes,
    unique_indices,
    indices_are_sorted,
    mode,
    fill_value,
):
  in_aval = ctx.avals_in[0]
  indices_aval = ctx.avals_in[1]
  out_aval = ctx.avals_out[0]

  if (
      len(in_aval.shape) != len(out_aval.shape)
      or indices_aval.shape[:-1] != out_aval.shape
      or indices_aval.shape[-1] != 1
  ):
    raise NotImplementedError("Only take_along_axis-like gathers supported")
  rank = len(out_aval.shape)

  # During lowering jnp.take_along_axis to lax.gather, we append extra dimension
  # to the end of the indices array. We should reshape it back to the original
  # shape before lowering to Mosaic and rely on MLIR canonicalization to remove
  # the reshapes.
  recovered_indices = vector.shape_cast(
      ir.VectorType.get(
          ctx.lowering_context.dynamic_shape_replacement_fn(out_aval.shape),
          indices.type.element_type,
      ),
      indices,
  )
  # Note: current support for lax.gather is still very limited.
  del fill_value

  (
      offset_dims,

View on GitHub (pinned to 1e1c6a8fc0)

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

Example fix

// before
out = lax.gather(x, idx_map, dimension_numbers=..., slice_sizes=...)
// after
out = jnp.take_along_axis(x, idx[..., None], axis=-1)

Defensive patterns

Strategy: validation

Validate before calling

import jax.numpy as jnp
def check_take_along_shape(x_shape, idx_shape, axis):
    out_shape = list(x_shape); out_shape.pop(axis)
    return len(idx_shape) == len(out_shape) + 1 and idx_shape[-1] == 1 and idx_shape[:-1] == tuple(out_shape)

Try / catch

catch NotImplementedError around kernel compilation and fall back to running the gather in plain JAX outside the kernel

Prevention

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

When it happens

Trigger: 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.

Common situations: 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.

Related errors

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/1defceae284a07c6. Report an issue: GitHub.