jax-ml/jax · error · NotImplementedError

Unsupported memory space when lowering memref.cast: {memory_

Error message

Unsupported memory space when lowering memref.cast: {memory_space}

What it means

memref.cast lowering in Mosaic only handles shared memory (smem) and tensor memory (tmem). A cast whose result memory space is anything else hits the terminal NotImplementedError.

Source

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

      )
    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)
def _memref_transpose_op_lowering_rule(
    ctx: LoweringContext, op: memref.TransposeOp
) -> Sequence[ir.Value]:
  del ctx

  in_transforms_attr = inference_utils.in_transforms(op)[0]
  unwrapped_in_ref = unwrap_transformed_memref(op.in_, in_transforms_attr)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Restrict casts to smem or tmem results; for global memory use plain loads/stores or load_tensor/store_tensor ops
2. If you need a global cast, lower it before Mosaic (standard MLIR memref.cast is a no-op) and skip the Mosaic rule
3. Check the memory_space attribute value on the result type and fix the type construction

def cast_memory_space_supported(ty) -> bool:
    ms = ir.MemRefType(ty).memory_space
    return ms == utils.smem() or ms == utils.tmem()

When it happens

Trigger: Constructing memref.cast where ir.MemRefType(op.result.type).memory_space is neither utils.smem() nor utils.tmem() (e.g. global/-generic memory space attribute).

Common situations: Hand-building Mosaic IR with custom memory space integer attributes, or a version change where a new memory space enum value isn't yet handled by the lowering.

Related errors

• Only support memref.cast where the input and output types ar
• memref.cast transforms must have identical transforms for bo
• Only SMEM and TMEM refs are supported.
• Only workgroup memory is supported but got {ref}.
• Unsupported memory space.

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