{"record":{"id":"9df1e62d61bca43f","repo":"jax-ml/jax","slug":"lhs-must-be-a-tmem-smem-ref","errorCode":null,"errorMessage":"LHS must be a TMEM/SMEM Ref.","messagePattern":"LHS must be a TMEM/SMEM Ref\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/primitives.py","lineNumber":2570,"sourceCode":"@tcgen05_mma_p.def_effectful_abstract_eval\ndef _tcgen05_mma_abstract_eval(acc, a, b, accumulate,\n                               *barrier_scales_and_transforms_leaves,\n                               acc_transforms_tree, a_transforms_tree,\n                               b_transforms_tree,\n                               barrier_transforms_tree,\n                               a_scale_transforms_tree,\n                               b_scale_transforms_tree,\n                               a_sparse_metadata_transforms_tree,\n                               collective_axis,\n                               arrive,\n                               scaled,\n                               sparse):\n  del accumulate, acc_transforms_tree, a_transforms_tree, b_transforms_tree, barrier_transforms_tree\n\n  if acc.memory_space != gpu_core.TMEM:\n    raise ValueError(\"Accumulator must be a TMEM Ref.\")\n  if a.memory_space not in (gpu_core.SMEM, gpu_core.TMEM):\n    raise ValueError(\"LHS must be a TMEM/SMEM Ref.\")\n  if b.memory_space != gpu_core.SMEM:\n    raise ValueError(\"RHS must be an SMEM Ref.\")\n\n  if collective_axis is not None:\n    # TODO(justinfu): If under a core_map, the avals for acc/a\n    # become normal MemRefs so we cannot check if they are collective.\n    # 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(","sourceCodeStart":2552,"sourceCodeEnd":2588,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/primitives.py#L2552-L2588","documentation":"The LHS operand of tcgen05.mma must be a ref in either SMEM or TMEM; a GMEM ref or other memory space fails this check.","triggerScenarios":"Passing a raw kernel input (GMEM) directly as the LHS without first loading/swizzling it into SMEM or TMEM.","commonSituations":"Forgetting to stage the A tile through an SMEM scratch buffer with the required swizzle/tiling; treating tcgen05.mma like a regular dot that accepts GMEM operands.","solutions":["Load the LHS tile into an SMEM buffer with proper swizzle before the MMA","Or keep A in TMEM if using a TMEM-resident LHS layout","Ensure the ref's memory_space attribute is SMEM or TMEM"],"exampleFix":"# before\ntcgen05.mma(a_gmem_ref, b_smem, acc, k_dim=k)\n# after\na_smem = load_to_smem(a_gmem_ref, swizzle=128)\ntcgen05.mma(a_smem, b_smem, acc, k_dim=k)","handlingStrategy":"type-guard","validationCode":"assert a.memory_space in (gpu_core.SMEM, gpu_core.TMEM)","typeGuard":"def lhs_space_ok(ref):\n    return getattr(ref, 'memory_space', None) in (gpu_core.SMEM, gpu_core.TMEM)","tryCatchPattern":null,"preventionTips":["Always stage operands into SMEM/TMEM scratch before tcgen05.mma","Write a kernel-setup helper validating all memory spaces"],"tags":["jax","pallas","tcgen05","smem","tmem","memory-space"],"backgroundTag":"gpu-memory-space-validation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}