jax-ml/jax · error · NotImplementedError

memref.cast tmem layouts must be identical for both input an

Error message

memref.cast tmem layouts must be identical for both input and output but got {in_tmem_layout=} and {out_tmem_layout=}

What it means

For tensor-memory (tmem) casts in Mosaic GPU, the input and output TMEM layout annotations must match exactly; the lowering just re-wraps the same tmem reference and cannot reconcile different layouts.

Source

Thrown at jax/experimental/mosaic/gpu/dialect_lowering.py:2102

  if memory_space == utils.smem():
    [in_transforms] = inference_utils.in_transforms(op)
    [out_transforms] = inference_utils.out_transforms(op)
    if in_transforms != out_transforms:
      raise NotImplementedError(
          "memref.cast transforms must have identical transforms for both "
          f"input and output but got {in_transforms=} and {out_transforms=}"
      )
    result = memref.cast(
        transform_type(ir.MemRefType(op.result.type), out_transforms),
        unwrap_transformed_memref(op.source, in_transforms),
    )
    return [wrap_transformed_memref(result, op.result.type, out_transforms)]

  if memory_space == utils.tmem():
    [in_tmem_layout] = inference_utils.in_tmem_layouts(op)
    [out_tmem_layout] = inference_utils.out_tmem_layouts(op)
    if in_tmem_layout != out_tmem_layout:
      raise NotImplementedError(
          "memref.cast tmem layouts must be identical for both input and"
          f" output but got {in_tmem_layout=} and {out_tmem_layout=}"
      )
    return [_tmem_ref_to_ir(_tmem_ref_from_ir(op.source, in_tmem_layout),
                            op.result.type)]

  raise NotImplementedError(
      f"Unsupported memory space when lowering memref.cast: {memory_space}"
  )


def _permutation_to_affine_map_attr(
    permutation: Sequence[int],
) -> ir.AffineMapAttr:
  return ir.AffineMapAttr.get(ir.AffineMap.get_permutation(permutation))


@_register_lowering(memref.TransposeOp, support_warp_semantics=True)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use the same TMEM layout annotation on both operands (pass the source layout through)
2. Allocate the destination with the matching layout instead of relayouting via cast
3. Print in/out tmem layouts via inference_utils.in_tmem_layouts(op) to find the divergence

[in_l] = inference_utils.in_tmem_layouts(op)
[out_l] = inference_utils.out_tmem_layouts(op)
assert in_l == out_l, 'tmem layouts must match for cast'

When it happens

Trigger: memref.cast with memory_space == utils.tmem() where the op's in_tmem_layouts and out_tmem_layouts annotations differ (different tmem packing/dp word layout).

Common situations: Working with tcgen05 tensor memory on Blackwell: copying a tmem ref and simultaneously trying to reinterpret its layout; mismatch usually comes from layout inference defaults when a cast result type was constructed by hand.

Related errors

• packed cannot be specified if layout is specified.
• packed, collective and layout arguments are only supported f
• layout attribute is only defined for TMEM refs
• collective attribute is only defined for TMEM refs
• Some aliased TMEM references are collective and some are not

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