jax-ml/jax · error · NotImplementedError

Transforms not supported for matmul_acc_lhs.

Error message

Transforms not supported for matmul_acc_lhs.

What it means

Raised during Mosaic/TPU lowering of a Pallas kernel that uses the tpu.matmul_acc_lhs primitive when accumulator transforms are supplied. The lowering path for this matmul variant simply does not implement transforms (e.g. scaling/offset accumulations), so it refuses to compile rather than producing wrong results.

Source

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

  mxu_id: int

  def __post_init__(self):
    tpu_core.check_accumulator_ref(self.shape, self.dtype, self.mxu_id)


@register_lowering_rule(tpu_primitives.matmul_acc_lhs_p)
def _matmul_acc_lhs_lowering_rule(
    ctx: LoweringRuleContext,
    acc: AccRef,
    lhs: ir.Value,
    *flat_acc_transforms,
    load_staged_rhs: int | None,
    acc_transforms_tree,
):
  del ctx
  acc_transforms = jax.tree.unflatten(acc_transforms_tree, flat_acc_transforms)
  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(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove accumulator transforms from the kernel using matmul_acc_lhs
2. Use the regular tpu.matmul primitive instead, which supports transforms
3. Check the JAX version release notes for matmul_acc_lhs transform support before upgrading

Example fix

// before
tpu.matmul_acc_lhs(acc, lhs, acc_transforms=transforms)
// after
tpu.matmul(lhs, rhs, acc=acc)  # no transforms on acc

from jax.tree_util import tree_flatten
flat, treedef = tree_flatten(acc_transforms)
assert not flat, 'matmul_acc_lhs does not support acc transforms'

When it happens

Trigger: Calling tpu.matmul_acc_lhs(...) (directly or via a Pallas TPU kernel) where the accumulator carries transform metadata — typically when jit transformations or accumulation transforms propagate into the primitive's lowering.

Common situations: Writing custom TPU Pallas matmul kernels against jax._src.pallas.mosaic.tpu and composing them with transformed accumulators; version upgrades where transforms started being threaded through matmul primitives.

Related errors

• Transforms not supported for matmul_pop.
• 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/c34ddb9728795853. Report an issue: GitHub.