{"record":{"id":"678b3ed3becebde9","repo":"jax-ml/jax","slug":"transforms-not-supported-for-matmul-pop","errorCode":null,"errorMessage":"Transforms not supported for matmul_pop.","messagePattern":"Transforms not supported for matmul_pop\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/lowering.py","lineNumber":5778,"sourceCode":"  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(\n      ir.VectorType.get(\n          ctx.lowering_context.dynamic_shape_replacement_fn(acc.shape),\n          _dtype_to_ir_type(acc.dtype)),\n      acc.base_address,\n      acc.mxu_id,\n  )\n\n\n@register_lowering_rule(tpu_primitives.matmul_lhs_fifo_p)\ndef _matmul_lhs_fifo_lowering_rule(\n    ctx: LoweringRuleContext,\n    lhs: ir.Value,\n    *,\n    mxu_index: int,\n    load_staged_rhs: int | None,\n):\n  del ctx","sourceCodeStart":5760,"sourceCodeEnd":5796,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/lowering.py#L5760-L5796","documentation":"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.","triggerScenarios":"Invoking tpu.matmul_pop(...) in a Pallas TPU kernel where acc_transforms is non-empty after tree unflattening.","commonSituations":"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.","solutions":["Clear accumulator transforms before calling matmul_pop","Restructure the kernel to avoid popping a transformed accumulator","Pin/choose a JAX version whose matmul_pop path supports your transform"],"exampleFix":"// before\ntpu.matmul_pop(acc, acc_transforms=transforms)\n// after\ntpu.matmul_pop(acc)  # transforms removed","handlingStrategy":"validation","validationCode":"flat, _ = tree_flatten(acc_transforms)\nassert not flat, 'matmul_pop does not support acc transforms'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Audit kernels for transformed accumulators before compilation","Prefer tpu.matmul for transform-heavy accumulation"],"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"}