jax-ml/jax · error · NotImplementedError

Unsupported gather

Error message

Unsupported gather

What it means

Catch-all failure in the Mosaic lax.gather lowering rule: the gather's dimension numbers did not match any of the recognized patterns (offset-only collapsed dims single-axis, or start_index_map gathers). The TPU backend only implements a narrow subset of JAX's general gather semantics.

Source

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

          lax.GatherScatterMode.FILL_OR_DROP,
          lax.GatherScatterMode.PROMISE_IN_BOUNDS,
      )
      and not offset_dims
      and collapsed_slice_dims == start_index_map
      and operand_batching_dims == start_indices_batching_dims
      and len(collapsed_slice_dims) == 1
      and len(operand_batching_dims) == rank - 1
  ):
    (axis,) = collapsed_slice_dims
    if (
        ctx.lowering_context.kernel_type == tpu_core.CoreType.TC
        and axis < rank - 2
    ):
      raise NotImplementedError(
          "Only gathers along the two minormost dimensions supported on TC"
      )
    return tpu.dynamic_gather(x, recovered_indices, [axis])
  raise NotImplementedError("Unsupported gather")


@register_lowering_rule(lax.transpose_p)
def _transpose_lowering_rule(ctx: LoweringRuleContext, x, *, permutation):
  out_type = ctx.aval_to_ir_type(ctx.avals_out[0])
  return tpu.transpose(out_type, x, permutation)


def _bcast(
    x: ir.Value | object,
    y: ir.Value | object,
    x_aval: ShapedAbstractValue,
    y_aval: ShapedAbstractValue,
    out_aval: ShapedAbstractValue,
    dynamic_shape_replacement_fn: DynamicShapeReplacementFn,
) -> tuple[ir.Value, ir.Value]:
  x_dtype = x_aval.dtype
  y_dtype = y_aval.dtype

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Re-express the operation as jnp.take_along_axis on a single axis
2. Compute the gather outside the Pallas kernel
3. Rewrite using explicit tpu primitives or manual index arithmetic with tpu.load
4. Upgrade JAX — gather coverage in Mosaic improves over versions

Example fix

// before
vals = lax.gather(x, index_cm, dnums, slice_sizes)
// after
vals = x[flat_idx]  # computed outside kernel, passed as input

except NotImplementedError: fall back to computing the gather outside the Pallas kernel with plain jnp indexing

When it happens

Trigger: Calling lax.gather (directly or via jnp.ix_/fancy indexing traced into a kernel) with collapsed_slice_dims or start_index_map combinations outside the two supported cases in the lowering rule.

Common situations: Using arbitrary fancy indexing inside jax.pallas kernels; mixing gather with batched index arrays the lowering doesn't recognize; relying on NumPy-style advanced indexing semantics in kernel code.

Related errors

• Only take_along_axis-like gathers supported
• masked load_p
• run_scoped_p with collective axes is not supported
• Non-decrementing wait is not supported.
• Only gathers along the two minormost dimensions supported on

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