{"record":{"id":"c34ddb9728795853","repo":"jax-ml/jax","slug":"transforms-not-supported-for-matmul-acc-lhs","errorCode":null,"errorMessage":"Transforms not supported for matmul_acc_lhs.","messagePattern":"Transforms not supported for matmul_acc_lhs\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/lowering.py","lineNumber":5761,"sourceCode":"  mxu_id: int\n\n  def __post_init__(self):\n    tpu_core.check_accumulator_ref(self.shape, self.dtype, self.mxu_id)\n\n\n@register_lowering_rule(tpu_primitives.matmul_acc_lhs_p)\ndef _matmul_acc_lhs_lowering_rule(\n    ctx: LoweringRuleContext,\n    acc: AccRef,\n    lhs: ir.Value,\n    *flat_acc_transforms,\n    load_staged_rhs: int | None,\n    acc_transforms_tree,\n):\n  del ctx\n  acc_transforms = jax.tree.unflatten(acc_transforms_tree, flat_acc_transforms)\n  if acc_transforms:\n    raise NotImplementedError(\"Transforms not supported for matmul_acc_lhs.\")\n  staged_rhs_kwarg: dict[str, Any] = {}\n  if load_staged_rhs is not None:\n    staged_rhs_kwarg = {\"load_staged_rhs\": load_staged_rhs}\n  tpu.matmul_acc_lhs(acc.base_address, lhs, acc.mxu_id, **staged_rhs_kwarg)\n  return []\n\n\n@register_lowering_rule(tpu_primitives.matmul_pop_p)\ndef _matmul_pop_lowering_rule(\n    ctx: LoweringRuleContext,\n    acc: AccRef,\n    *flat_acc_transforms,\n    acc_transforms_tree,\n):\n  acc_transforms = jax.tree.unflatten(acc_transforms_tree, flat_acc_transforms)\n  if acc_transforms:\n    raise NotImplementedError(\"Transforms not supported for matmul_pop.\")\n  return tpu.matmul_pop(","sourceCodeStart":5743,"sourceCodeEnd":5779,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/lowering.py#L5743-L5779","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Remove accumulator transforms from the kernel using matmul_acc_lhs","Use the regular tpu.matmul primitive instead, which supports transforms","Check the JAX version release notes for matmul_acc_lhs transform support before upgrading"],"exampleFix":"// before\ntpu.matmul_acc_lhs(acc, lhs, acc_transforms=transforms)\n// after\ntpu.matmul(lhs, rhs, acc=acc)  # no transforms on acc","handlingStrategy":"validation","validationCode":"from jax.tree_util import tree_flatten\nflat, treedef = tree_flatten(acc_transforms)\nassert not flat, 'matmul_acc_lhs does not support acc transforms'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Keep accumulator transforms off matmul_acc_lhs/matmul_pop kernels","Pin supported JAX versions when using exotic matmul primitives"],"tags":["jax","pallas","tpu","matmul","not-implemented"],"backgroundTag":"unsupported-operation-not-implemented","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}