{"record":{"id":"ed30b1104f43d7b6","repo":"jax-ml/jax","slug":"lhs-ref-must-be-collective-if-collective-axis-is-s","errorCode":null,"errorMessage":"LHS Ref must be collective if collective_axis is set.","messagePattern":"LHS Ref must be collective if collective_axis is set\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/primitives.py","lineNumber":2582,"sourceCode":"                               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(\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","sourceCodeStart":2564,"sourceCodeEnd":2600,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/primitives.py#L2564-L2600","documentation":"When collective_axis is set on tcgen05.mma and the LHS is a TMEM ref, that TMEM ref must be marked collective; otherwise the LHS layout does not match the 2CTA tensor-core expectation.","triggerScenarios":"Passing a non-collective TMEM-resident LHS together with collective_axis in tcgen05.mma.","commonSituations":"Using TMEM for A in a kernel later converted to collective MMA without updating A's allocation; mixing collective accumulator with per-CTA LHS TMEM.","solutions":["Mark the LHS TMEM ref as collective when using collective_axis","Or move the LHS to SMEM (SMEM path is not subject to this check)","Keep the collective flag consistent across acc/a allocations"],"exampleFix":"# before\na_tmem = allocate(TMEM, (m, k), jnp.float8e4m3)\ntcgen05.mma(a_tmem, b_smem, acc, k_dim=k, collective_axis=0)\n# after\na_tmem = allocate(TMEM, (m, k), jnp.float8e4m3, collective=True)\ntcgen05.mma(a_tmem, b_smem, acc, k_dim=k, collective_axis=0)","handlingStrategy":"validation","validationCode":"if collective_axis is not None and isinstance(a, gpu_core.AbstractTMEMRef):\n    assert a.collective, 'LHS TMEM must be collective'","typeGuard":"def lhs_collective_ok(a, collective_axis):\n    if collective_axis is None:\n        return True\n    return not isinstance(a, gpu_core.AbstractTMEMRef) or a.collective","tryCatchPattern":null,"preventionTips":["Keep collective flags of acc and LHS in sync","Prefer SMEM LHS when unsure about collective TMEM layouts"],"tags":["jax","pallas","tcgen05","collective-mma","tmem"],"backgroundTag":"gpu-memory-space-validation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}