jax-ml/jax · error · NotImplementedError

memref.cast transforms must have identical transforms for bo

Error message

memref.cast transforms must have identical transforms for both input and output but got {in_transforms=} and {out_transforms=}

What it means

When casting a shared-memory memref, Mosaic requires the transform annotations on input and output to be exactly identical, because smem casts are pure pointer reinterpretations. Any divergence in tiling/swizzle annotations is rejected.

Source

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

  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)]

  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),

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Copy the source's transform annotation verbatim to the cast result
Use a dedicated transform op (tile/swizzle) to change layout rather than cast
Inspect both transform attrs with inference_utils.in/out_transforms(op) and diff them

Example fix

// before
out = t.memref.cast(in_smem, ty_with_different_transforms)
// after
[in_t] = in_transforms  # reuse
out = t.memref.cast(in_smem, transform_type(ty, in_t))

Defensive patterns

Strategy: validation

Validate before calling

[in_t] = inference_utils.in_transforms(op)
[out_t] = inference_utils.out_transforms(op)
assert in_t == out_t, 'smem cast transforms must match'

Type guard

def smem_cast_transforms_match(op) -> bool:
    return inference_utils.in_transforms(op) == inference_utils.out_transforms(op)

Prevention

Copy source transform annotation to cast results
Avoid changing layout via cast

When it happens

Trigger: memref.cast into utils.smem() where inference_utils.in_transforms(op) != out_transforms(op), e.g. annotating the result with a different tile shape or swizzle than the source.

Common situations: Manually setting transforms on a cast result while the source carries inferred transforms; or re-tiling a tensor in smem via cast instead of an explicit transform op.

Related errors

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