jax-ml/jax · error · NotImplementedError

SubViewOp only supports a single tile transform.

Error message

SubViewOp only supports a single tile transform.

What it means

Mosaic GPU's lowering of memref.subview only handles a single tile transform on the source memref. When the input has zero or multiple transforms, the pattern match falls into the default case and lowering fails.

Source

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

        )
      new_sizes = tile_transform.transform_shape(list(op.static_sizes))
      # TODO(bchetioui): support transposed offsets.
      new_static_offsets, new_dynamic_offsets = _tile_transform_offsets(
          tiling, list(op.static_offsets), list(op.offsets)
      )

      new_subview_op = memref.SubViewOp(
          transform_type(ir.MemRefType(op.result.type), transforms),
          unwrapped_source_ref,
          new_dynamic_offsets,
          sizes=[],
          strides=[],
          static_offsets=new_static_offsets,
          static_sizes=new_sizes,
          static_strides=[1] * len(in_transformed_ty.shape),
      )
    case _:
      raise NotImplementedError(
          "SubViewOp only supports a single tile transform."
      )

  wrapped_ref = wrap_transformed_memref(
      new_subview_op.result, op.result.type, out_transforms
  )
  return [wrapped_ref]


# memref.cast shows up when we slice a ref with a dynamic index, that later gets
# folded into a constant. At that time, the sliced ref type is simplified from
# having a dynamic offset to a constant offset. However, downstream consumer ops
# still expect the offset to be dynamic, forcing the insertion of a memref.cast
# op to reintroduce the dynamic offset.
@_register_lowering(memref.CastOp, support_warp_semantics=True)
def _memref_cast_op_lowering_rule(
    ctx: LoweringContext, op: memref.CastOp
) -> Sequence[ir.Value]:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Reduce the memref to a single tile transform before taking a subview (apply the other transform manually via explicit reshape/layout ops)
2. Check inference_utils.in_transforms(op) in a debug pass to see why multiple transforms were inferred
3. Avoid slicing transformed memrefs; instead slice first and then apply the tiling transform

Example fix

// before
sliced = t.memref_subview(tiled_and_swizzled_ref, offsets, sizes)
// after
plain = t.memref_subview(plain_ref, offsets, sizes)
tiled = t.tile(plain, tile_shape)  # single transform

in_ts = inference_utils.in_transforms(op)
assert len(in_ts) <= 1, f'subview needs <=1 transform, got {len(in_ts)}'

def has_single_tile_transform(op) -> bool:
    ts = inference_utils.in_transforms(op)
    return len(ts) == 1 and isinstance(ts[0], lc.TileTransform)

When it happens

Trigger: Calling t.memref_subview (or building a SubViewOp) on a memref whose in_transforms contains more than one MemRefTransform (e.g. nested/stacked TileTransforms) during the Mosaic-to-MemRef dialect lowering pass.

Common situations: Composing multiple layout transforms (e.g. swizzle + tiling, or two tilings) on a shared-memory tensor and then slicing it in a Mosaic GPU kernel; typically after upgrading JAX where transform inference became stricter.

Related errors

• Only support memref.cast where the input and output types ar
• memref.cast transforms must have identical transforms for bo
• Unsupported memory space when lowering memref.cast: {memory_
• Unsupported in/out transforms. In transform: {in_transform},
• Only a single tiling transform is supported when collapsing

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