{"record":{"id":"75fe1ac979ab9bff","repo":"jax-ml/jax","slug":"b-scale-must-be-a-tmem-ref","errorCode":null,"errorMessage":"b_scale must be a TMEM Ref","messagePattern":"b_scale must be a TMEM Ref","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/primitives.py","lineNumber":2597,"sourceCode":"      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,\n    barrier_transforms_tree,\n    a_scale_transforms_tree,","sourceCodeStart":2579,"sourceCodeEnd":2615,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/primitives.py#L2579-L2615","documentation":"In a block-scaled tcgen05.mma, the b_scale operand must be a TMEM ref, mirroring the a_scale requirement.","triggerScenarios":"Passing b_scale in SMEM/GMEM to a scaled tcgen05.mma.","commonSituations":"Same as a_scale: scales staged in SMEM alongside B instead of TMEM.","solutions":["Allocate b_scale in TMEM","Keep both scale refs' memory spaces consistent (TMEM)","Add a pre-call assert on memory_space for both scales"],"exampleFix":"# before\nb_scale = allocate(SMEM, scale_shape, jnp.uint8)\n# after\nb_scale = allocate(TMEM, scale_shape, jnp.uint8)","handlingStrategy":"type-guard","validationCode":"if scaled:\n    assert b_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","Add a shared assert for a_scale and b_scale"],"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"}