{"record":{"id":"6586b591b816d18e","repo":"jax-ml/jax","slug":"a-scale-must-be-a-tmem-ref","errorCode":null,"errorMessage":"a_scale must be a TMEM Ref","messagePattern":"a_scale must be a TMEM Ref","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/primitives.py","lineNumber":2595,"sourceCode":"    # Figure out a way to fix this.\n    if isinstance(acc, gpu_core.AbstractTMEMRef) and not acc.collective:\n      raise ValueError(\n          \"Accumulator Ref must be collective if collective_axis is set.\")\n    if isinstance(a, gpu_core.AbstractTMEMRef) and not a.collective:\n      raise ValueError(\n          \"LHS Ref must be collective if collective_axis is set.\")\n\n  scales_and_transforms_leaves = barrier_scales_and_transforms_leaves\n  if arrive:\n    barrier, *scales_and_transforms_leaves = barrier_scales_and_transforms_leaves\n    orders_tensor_core = getattr(\n        barrier.inner_aval.dtype, \"orders_tensor_core\", False)\n    if not orders_tensor_core:\n      raise ValueError(\"MMA barrier must have orders_tensor_core set to True.\")\n  if scaled:\n    a_scale, b_scale = scales_and_transforms_leaves[:2]\n    if a_scale.memory_space != gpu_core.TMEM:\n      raise ValueError(\"a_scale must be a TMEM Ref\")\n    if b_scale.memory_space != gpu_core.TMEM:\n      raise ValueError(\"b_scale must be a TMEM Ref\")\n\n  return [], {gpu_core._memory_effect}\n\n\n@lowering.register_lowering_rule(tcgen05_mma_p, *gpu_core.LANExWG_SEMANTICS)\n@lowering.register_lowering_rule(tcgen05_mma_p, *gpu_core.LANExWARP_SEMANTICS)\ndef _tcgen05_mma_lowering(\n    ctx: lowering.LoweringRuleContext,\n    acc: tcgen05.TMEMRef,\n    a_ref,\n    b_ref,\n    accumulate: bool | ir.Value,\n    *barrier_scales_and_transforms_leaves,\n    acc_transforms_tree,\n    a_transforms_tree,\n    b_transforms_tree,","sourceCodeStart":2577,"sourceCodeEnd":2613,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/primitives.py#L2577-L2613","documentation":"In a block-scaled tcgen05.mma (scaled=True), the a_scale operand must be a TMEM ref; SMEM/GMEM scale refs are rejected.","triggerScenarios":"Passing a_scale allocated in SMEM or as a raw GMEM input to a scaled tcgen05.mma call.","commonSituations":"Allocating scales next to the operand tiles in SMEM instead of TMEM when implementing MXFP8 block scaling.","solutions":["Allocate a_scale in TMEM","Follow the scaled-MMA example's scale allocation pattern exactly","Verify a_scale.memory_space == TMEM before calling"],"exampleFix":"# before\na_scale = allocate(SMEM, scale_shape, jnp.uint8)\ntcgen05.mma(a, b, acc, k_dim=k, a_scale=a_scale, b_scale=b_scale)\n# after\na_scale = allocate(TMEM, scale_shape, jnp.uint8)\ntcgen05.mma(a, b, acc, k_dim=k, a_scale=a_scale, b_scale=b_scale)","handlingStrategy":"type-guard","validationCode":"if scaled:\n    assert a_scale.memory_space == gpu_core.TMEM","typeGuard":"def scale_is_tmem(ref):\n    return getattr(ref, 'memory_space', None) == gpu_core.TMEM","tryCatchPattern":null,"preventionTips":["Allocate both scales in TMEM in one place","Mirror the scaled-MMA example's allocation code"],"tags":["jax","pallas","tcgen05","block-scaling","tmem"],"backgroundTag":"gpu-memory-space-validation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}