jax-ml/jax · error · NotImplementedError

Unsupported TMEM ref {ref}.

Error message

Unsupported TMEM ref {ref}.

What it means

For TMEM aliases, the lowering can only recover the base offset if the ref's owner chain is a SliceTmemOp (optionally behind a TmemLayoutCastOp). Any other op producing the TMEM ref (arbitrary layout casts, allocs, other dialect ops) makes offset computation impossible and NotImplementedError is raised.

Source

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

          else:
            ref_bytes = ref_bits // 8
            ref = mgpu.memref_slice(ref, slice(offset, offset + ref_bytes))
            ref = _handle_dtype_bitcast(
                ref,
                ir.MemRefType(ref.type).element_type,
                mlir_dtype,
            )
            ref = mgpu.memref_reshape(ref, transformed_shape)
        elif input_ref_ty.memory_space == mgpu_utils.tmem():

          if isinstance(ref.owner, mgpu.dialect.SliceTmemOp):
            source_slice_op = ref.owner
          elif isinstance(
              ref.owner, mgpu.dialect.TmemLayoutCastOp
          ) and isinstance(ref.owner.operands[0].owner, mgpu.dialect.SliceTmemOp):
            source_slice_op = ref.owner.operands[0].owner
          else:
            raise NotImplementedError(f"Unsupported TMEM ref {ref}.")

          base_offset = source_slice_op.offset.value
          assert isinstance(base_offset, int)  # make pyrefly happy
          total_offset = base_offset + offset
          ref_ty = ir.MemRefType.get(
              transformed_shape, mlir_dtype, memory_space=mgpu_utils.tmem()
          )
          alloc_id = source_slice_op.alias_id
          assert alloc_id is not None
          # TODO(bchetioui): Use a scheme resilient to hash collisions.
          alias_id = hash((offset, alloc_id.value, alias_group_idx))
          slice_op = mgpu.dialect.SliceTmemOp(
              ref_ty, source_slice_op.source, total_offset
          )
          i64 = ir.IntegerType.get_signless(64)
          slice_op.attributes["alias_id"] = ir.IntegerAttr.get(i64, alias_id)
          ref = slice_op.result
          assert layout is not None

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Ensure the aliased TMEM ref comes directly from a tmem slice op (slice before any other transform)
2. Apply the alias view earlier in the op chain, immediately after slicing
3. Simplify/reorder TMEM layout casts so a slice_tmem is the immediate (or one-hop) owner
4. Upgrade JAX / report upstream with a minimal kernel repro

Example fix

# before
r = tmem_alloc(...)            # owner not a slice op
v = r.view(dtype)
# after
r = tmem_alloc(...)
r = r[0:n]                     # slice_tmem owner
v = r.view(dtype)

owner = getattr(ref, 'owner', None)
ok = type(owner).__name__ == 'SliceTmemOp' or (
     type(owner).__name__ == 'TmemLayoutCastOp' and
     type(getattr(owner.operands[0], 'owner', None)).__name__ == 'SliceTmemOp')
assert ok, 'TMEM alias must come from slice_tmem (optionally behind tmem_layout_cast)'

def tmem_ref_aliasable(ref) -> bool:
    o = getattr(ref, 'owner', None)
    if type(o).__name__ == 'SliceTmemOp':
        return True
    if type(o).__name__ == 'TmemLayoutCastOp':
        return type(getattr(o.operands[0], 'owner', None)).__name__ == 'SliceTmemOp'
    return False

When it happens

Trigger: Aliasing a TMEM ref whose owner is not slice_tmem / tmem_layout_cast(slice_tmem) — e.g. aliasing a raw TMEM allocation or the result of another TMEM-transforming op.

Common situations: Blackwell tcgen05 kernels chaining multiple TMEM layout operations before an aliased view; evolving JAX versions adding new TMEM ops the alias path doesn't recognize.

Related errors

• TMEM aliasing only supported for Refs with the same first di
• Some aliased TMEM references are collective and some are not
• All aliased Refs must have the same memory space (SMEM or TM
• Sparse metadata format not implemented for {operand_dtype=}
• The base ref for aliases must come from a slice_smem op.

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