jax-ml/jax · error · NotImplementedError

Transforms not supported for matmul_pop.

Error message

Transforms not supported for matmul_pop.

What it means

The Mosaic lowering of tpu.matmul_pop rejects any accumulator transforms, because that code path has no implementation for them. It mirrors the matmul_acc_lhs restriction: transform-free popping of the accumulator works, transformed accumulators do not compile.

Source

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

  if acc_transforms:
    raise NotImplementedError("Transforms not supported for matmul_acc_lhs.")
  staged_rhs_kwarg: dict[str, Any] = {}
  if load_staged_rhs is not None:
    staged_rhs_kwarg = {"load_staged_rhs": load_staged_rhs}
  tpu.matmul_acc_lhs(acc.base_address, lhs, acc.mxu_id, **staged_rhs_kwarg)
  return []


@register_lowering_rule(tpu_primitives.matmul_pop_p)
def _matmul_pop_lowering_rule(
    ctx: LoweringRuleContext,
    acc: AccRef,
    *flat_acc_transforms,
    acc_transforms_tree,
):
  acc_transforms = jax.tree.unflatten(acc_transforms_tree, flat_acc_transforms)
  if acc_transforms:
    raise NotImplementedError("Transforms not supported for matmul_pop.")
  return tpu.matmul_pop(
      ir.VectorType.get(
          ctx.lowering_context.dynamic_shape_replacement_fn(acc.shape),
          _dtype_to_ir_type(acc.dtype)),
      acc.base_address,
      acc.mxu_id,
  )


@register_lowering_rule(tpu_primitives.matmul_lhs_fifo_p)
def _matmul_lhs_fifo_lowering_rule(
    ctx: LoweringRuleContext,
    lhs: ir.Value,
    *,
    mxu_index: int,
    load_staged_rhs: int | None,
):
  del ctx

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Clear accumulator transforms before calling matmul_pop
2. Restructure the kernel to avoid popping a transformed accumulator
3. Pin/choose a JAX version whose matmul_pop path supports your transform

Example fix

// before
tpu.matmul_pop(acc, acc_transforms=transforms)
// after
tpu.matmul_pop(acc)  # transforms removed

flat, _ = tree_flatten(acc_transforms)
assert not flat, 'matmul_pop does not support acc transforms'

When it happens

Trigger: Invoking tpu.matmul_pop(...) in a Pallas TPU kernel where acc_transforms is non-empty after tree unflattening.

Common situations: Custom TPU kernels that pop MXU accumulators after a matmul while jit or accumulation transforms are active; upgrading JAX where the primitive signature gained acc_transforms.

Related errors

• Transforms not supported for matmul_acc_lhs.
• masked load_p
• run_scoped_p with collective axes is not supported
• Non-decrementing wait is not supported.
• Loading from an accumulator is not supported. Use `matmul_po

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