What causes "Unsupported copy: {src.type} -> {dst.type}" in jax-ml/jax?

Route SMEM-to-SMEM or GMEM-to-GMEM copies through registers instead: load into a FragmentedArray and store to the destination (or via an intermediate GMEM buffer).

jax-ml/jax · error · NotImplementedError

Unsupported copy: {src.type} -> {dst.type}

Error message

Unsupported copy: {src.type} -> {dst.type}

What it means

Raised by Mosaic GPU's low-level copy helper when asked to copy between two memrefs and exactly one of them is a shared-memory (SMEM) reference with 2D tiling; this path only supports SMEM<->GMEM copies. If both refs are SMEM, both are GMEM, or ranks/tiling don't match the expected pattern, there is no lowering implemented and the compiler raises NotImplementedError. It is a limitation of the Mosaic tiled-copy codegen, not a user data error.

Source

Thrown at jax/experimental/mosaic/gpu/fragmented_array.py:5304

          f"For {swizzle=}, expected SMEM tiling to be (8, {swizzle_elems})"
      )
    expected_src_shape = utils.tile_shape(gmem_ty.shape, (8, swizzle_elems))
    if tuple(smem_ty.shape) != expected_src_shape:
      raise ValueError(
          f"Expected SMEM reference to have shape {expected_src_shape} (tiling"
          f" {gmem_ty.shape} by (8, {swizzle_elems})), but got {smem_ty.shape}"
      )
    layout = tiled_copy_smem_gmem_layout(
        *smem_ty.shape[-4:-2], swizzle, bitwidth  # pyrefly: ignore[bad-argument-count]
    )
    if utils.is_smem_ref(src_ty):
      regs = FragmentedArray.load_tiled(src, swizzle, is_signed=is_signed, layout=layout)
      regs.store_untiled(dst, optimized=False)
    else:
      regs = FragmentedArray.load_untiled(src, is_signed=is_signed, layout=layout, optimized=False)
      regs.store_tiled(dst, swizzle)
    return
  raise NotImplementedError(f"Unsupported copy: {src.type} -> {dst.type}")


def is_supported_strided_layout_broadcast(
    src: WGStridedFragLayout,
    dst: WGStridedFragLayout,
    dims: tuple[int, ...],
) -> bool:
  """We only support broadcasting of leading dimensions."""
  if src.vec_size != dst.vec_size:
    return False
  # Check if input maps exactly to the end (prevents trailing dims).
  if dims != tuple(range(len(dst.shape) - len(src.shape), len(dst.shape))):
    return False
  # Identify input indices that are expanded vs. those that are preserved
  # Expansion: input is 1, output is > 1.
  # Preserved: input is > 1.
  exp_indices, pre_indices = [], []
  for i, dim in enumerate(src.shape):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Route SMEM-to-SMEM or GMEM-to-GMEM copies through registers instead: load into a FragmentedArray and store to the destination (or via an intermediate GMEM buffer).
Check utils.is_smem_ref on both operands before copying and restructure the kernel so copies cross SMEM/GMEM boundaries.
Verify the SMEM ref has the required rank+2 and trailing (8, 8*swizzle//bitwidth) tiling so the supported branch is taken.
Raise a feature request / check newer JAX for added copy combinations.

Example fix

// before
copy(src_smem_ref, dst_smem_ref, swizzle)  # both SMEM -> NotImplementedError
// after
regs = FragmentedArray.load_tiled(src_smem_ref, swizzle)
regs.store_tiled(dst_smem_ref, swizzle)

Defensive patterns

Strategy: validation

Validate before calling

from jax.experimental.mosaic.gpu import utils

def is_supported_copy(src, dst) -> bool:
    src_smem, dst_smem = utils.is_smem_ref(src.type), utils.is_smem_ref(dst.type)
    if src_smem == dst_smem:
        return False  # only SMEM<->GMEM pairs are lowered
    smem_ty, gmem_ty = (src.type, dst.type) if src_smem else (dst.type, src.type)
    return smem_ty.rank == gmem_ty.rank + 2

Try / catch

try:
    copy(src, dst, swizzle)
except NotImplementedError:
    regs = FragmentedArray.load_tiled(src, swizzle)
    regs.store_untiled(dst, optimized=False)

Prevention

Only issue tiled copies between one SMEM ref and one GMEM ref.
Validate is_smem_ref on both operands before copying.
Keep SMEM refs rank+2 with trailing (8, swizzle_elems) tiles.

When it happens

Trigger: Calling the Mosaic copy utility (e.g. via mosaicGPU copy ops / kernel code that copies between references) where src and dst are both in shared memory or both in global memory, or where is_smem_ref(src) == is_smem_ref(dst). Also triggered when a tiled SMEM copy is requested on refs that are not a (8, swizzle_elems)-tiled pair.

Common situations: Writing HSMEM-to-HSMEM or GMEM-to-GMEM copies in a Mosaic GPU kernel and expecting the tiled copy path to handle them; upgrading JAX versions where copy support was narrowed; passing incorrectly tiled SMEM refs (wrong swizzle/shape) so the code falls through to the final raise.

Related errors

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