jax-ml/jax · error · NotImplementedError

Only gathers along the two minormost dimensions supported on

Error message

Only gathers along the two minormost dimensions supported on TC

What it means

On TPU TensorCore cores, Mosaic can only lower dynamic gathers that operate on one of the two minormost (last two) dimensions of the tensor. A gather whose collapsed axis is further from the minor end falls back to an unsupported path on TC hardware and raises this NotImplementedError during kernel compilation.

Source

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

  if (
      slice_sizes == (1,) * rank
      and mode
      in (
          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,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Transpose the tensor so the gather axis becomes one of the last two dimensions, gather, then transpose back
Rework the kernel layout (e.g. use BlockMapping to put the gathered axis minormost)
Do the gather outside the kernel as a pre-processing step
Target a non-TC core type if the kernel doesn't need TensorCore ops

Example fix

// before
out = jnp.take_along_axis(x, idx, axis=0)  # x is 4D, axis 0 not minormost
// after
x_t = jnp.transpose(x, (1,2,3,0))
out_t = jnp.take_along_axis(x_t, idx_t, axis=-1)
out = jnp.transpose(out_t, (3,0,1,2))

Defensive patterns

Strategy: validation

Validate before calling

def gather_ok_on_tc(rank, axis):
    return axis >= rank - 2

Try / catch

catch NotImplementedError at compile time and restructure to transpose the gather axis minormost

Prevention

Keep gather axes minormost in kernel memory layouts
Document per-core-type assumptions in kernel docstrings
Test on the target core type in CI

When it happens

Trigger: Running a Pallas kernel on a TensorCore-targeted compilation where take_along_axis/dynamic_gather is along an axis with index < rank-2 (e.g. axis 0 of a 4D tensor).

Common situations: Writing a kernel that gathers along batch or heads dimension of a 4D activation tensor; code that worked on TPU v4/v5e scalar/vector paths failing when scheduled onto TC subcores; attention-style index selects on non-minor axes.

Related errors

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