jax-ml/jax · error · ValueError

TMEM aliasing only supported for Refs with the same first di

Error message

TMEM aliasing only supported for Refs with the same first dimension, got {ref.shape[0]} != {transformed_shape[0]}.

What it means

When aliasing tensor-memory (TMEM) refs, the lowering adjusts the base address by an offset while keeping the layout, which is only valid if the first dimension (the TMEM address/row dimension) is unchanged between the original ref and the transformed block. If ref.shape[0] != transformed_shape[0], aliasing would point at wrong rows.

Source

Thrown at jax/_src/pallas/mosaic_gpu/lowering.py:1573

) -> tuple[
    ir.Value | tcgen05.TMEMRef,
    state_types.AbstractRef,
    Sequence[state_types.Transform],
    Sequence[state_types.Transform],
]:
  # Looks for the first transform being an ExtractAliasedRef and pulls out the
  # Ref there, updating the transforms.
  match transforms:
    case (
        gpu_core.ExtractAliasedRef(dtype, transformed_shape, offset, alias_group_idx, layout) as t,
        *other_transforms,
    ):
      ref_aval = t.transform_type(ref_aval)
      mlir_dtype = mgpu_utils.dtype_to_ir_type(dtype)
      if isinstance(ref, tcgen05.TMEMRef):
        assert layout is not None
        if ref.shape[0] != transformed_shape[0]:
          raise ValueError(
              "TMEM aliasing only supported for Refs with the same first"
              f" dimension, got {ref.shape[0]} != {transformed_shape[0]}."
          )
        address = arith_dialect.addi(ref.address, _i32_constant(offset))
        ref = tcgen05.TMEMRef(
            address=address,
            shape=cast(tuple[int, int], transformed_shape),
            dtype=mlir_dtype,
            layout=layout,
        )
      else:
        assert isinstance(ref, ir.Value)  # make pyrefly happy
        input_ref_ty = ir.MemRefType(ref.type)
        if input_ref_ty.memory_space == mgpu_utils.smem():
          assert layout is None
          ref_bits = math.prod(transformed_shape) * mgpu_utils.bitwidth(
              mlir_dtype
          )

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Restructure the kernel so TMEM aliased refs keep the same leading dimension (apply the transform to other axes)
2. Move data to SMEM before applying shape-changing transforms
3. Split the kernel so the transform happens in a separate non-aliased step

Example fix

# before
@plt_kernel  # BlockMapping transforms dim 0 of tmem ref
# in_kernel: out_ref[...] aliases tmem_ref with transformed first dim
# after: keep dim0 fixed, transform trailing dims
out = tmem_ref[:, :k].view(...)  # dim0 unchanged

assert aliased.shape[0] == transformed.shape[0], 'TMEM alias must preserve the first dimension'

def tmem_alias_ok(ref_shape, transformed_shape) -> bool:
    return ref_shape[0] == transformed_shape[0]

When it happens

Trigger: A Pallas kernel using TMEM (tcgen05) refs where a BlockMapping or transformation changes the first dimension of the block — e.g. transpose-like or reshaping transforms applied to TMEM aliased refs.

Common situations: Blackwell tcgen05 kernels with aliased accumulator buffers under block mappings that reshape the leading axis; migrating SMEM aliasing patterns to TMEM.

Related errors

• Unsupported TMEM ref {ref}.
• Some aliased TMEM references are collective and some are not
• All aliased Refs must have the same memory space (SMEM or TM
• Stores to TMEM are asynchronous operations and cannot be per
• Accumulator and LHS have incompatible shapes. Expected LHS t

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/3b57c512b45dd73b. Report an issue: GitHub.