jax-ml/jax · error · NotImplementedError

Only support memref.cast where the input and output types ar

Error message

Only support memref.cast where the input and output types are the same up to offset, but got {in_ty=} and {out_ty=}.

What it means

memref.cast lowering requires the input and output memref types to be identical after erasing offsets. This error means the cast changes shape, strides, layout, or element type beyond a constant offset, which the lowering cannot represent.

Source

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

  out_ty = ir.MemRefType(op.result.type)
  in_strides, _ = in_ty.get_strides_and_offset()
  out_strides, _ = out_ty.get_strides_and_offset()

  unoffseted_in_ty = ir.MemRefType.get(
      in_ty.shape,
      in_ty.element_type,
      memory_space=in_ty.memory_space,
      layout=ir.StridedLayoutAttr.get(0, in_strides),
  )
  unoffseted_out_ty = ir.MemRefType.get(
      out_ty.shape,
      out_ty.element_type,
      memory_space=out_ty.memory_space,
      layout=ir.StridedLayoutAttr.get(0, out_strides),
  )

  if unoffseted_in_ty != unoffseted_out_ty:
    raise NotImplementedError(
        "Only support memref.cast where the input and output types are the "
        f"same up to offset, but got {in_ty=} and {out_ty=}."
    )

  memory_space = ir.MemRefType(op.result.type).memory_space
  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)]

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Make the cast type-preserving: only the offset may differ; keep shape, strides, element type and memory space equal
2. Use memref.reshape / expand_shape / collapse_shape for shape changes and a load+store or convert op for dtype changes
3. Verify with ir.MemRefType(x).get_strides_and_offset() that only the offset differs before emitting the cast

Example fix

// before
ref2 = t.memref.cast(ref, new_ty_with_different_shape)
// after
ref2 = t.memref.collapse_shape(ref, reassociation)  # shape change
ref3 = t.memref.cast(ref2, same_shape_offset_ty)     # cast only

in_ty, out_ty = ir.MemRefType(src.type), ir.MemRefType(dst.type)
assert in_ty.shape == out_ty.shape and in_ty.element_type == out_ty.element_type

def castable_up_to_offset(in_ty, out_ty) -> bool:
    return (in_ty.shape == out_ty.shape
            and in_ty.element_type == out_ty.element_type
            and in_ty.memory_space == out_ty.memory_space)

When it happens

Trigger: Issuing memref.cast where ir.MemRefType(source) and ir.MemRefType(result) differ in shape/strides/element type/memory space after normalizing offset to zero, e.g. casting a 2D memref to one with different dimensions.

Common situations: Using Mosaic's cast_to_smem/tmem helpers on a tensor whose layout was reshaped or re-strided first; mixing memory spaces or dtypes in a cast instead of using proper load/store conversion.

Related errors

• memref.cast transforms must have identical transforms for bo
• Unsupported memory space when lowering memref.cast: {memory_
• Expected an index-typed index
• SubViewOp only supports a single tile transform.
• memref.cast tmem layouts must be identical for both input an

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