{"record":{"id":"865c5f8382b8fede","repo":"jax-ml/jax","slug":"unsupported-copy-src-type-dst-type","errorCode":null,"errorMessage":"Unsupported copy: {src.type} -> {dst.type}","messagePattern":"Unsupported copy: (.+?) -> (.+?)","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/fragmented_array.py","lineNumber":5304,"sourceCode":"          f\"For {swizzle=}, expected SMEM tiling to be (8, {swizzle_elems})\"\n      )\n    expected_src_shape = utils.tile_shape(gmem_ty.shape, (8, swizzle_elems))\n    if tuple(smem_ty.shape) != expected_src_shape:\n      raise ValueError(\n          f\"Expected SMEM reference to have shape {expected_src_shape} (tiling\"\n          f\" {gmem_ty.shape} by (8, {swizzle_elems})), but got {smem_ty.shape}\"\n      )\n    layout = tiled_copy_smem_gmem_layout(\n        *smem_ty.shape[-4:-2], swizzle, bitwidth  # pyrefly: ignore[bad-argument-count]\n    )\n    if utils.is_smem_ref(src_ty):\n      regs = FragmentedArray.load_tiled(src, swizzle, is_signed=is_signed, layout=layout)\n      regs.store_untiled(dst, optimized=False)\n    else:\n      regs = FragmentedArray.load_untiled(src, is_signed=is_signed, layout=layout, optimized=False)\n      regs.store_tiled(dst, swizzle)\n    return\n  raise NotImplementedError(f\"Unsupported copy: {src.type} -> {dst.type}\")\n\n\ndef is_supported_strided_layout_broadcast(\n    src: WGStridedFragLayout,\n    dst: WGStridedFragLayout,\n    dims: tuple[int, ...],\n) -> bool:\n  \"\"\"We only support broadcasting of leading dimensions.\"\"\"\n  if src.vec_size != dst.vec_size:\n    return False\n  # Check if input maps exactly to the end (prevents trailing dims).\n  if dims != tuple(range(len(dst.shape) - len(src.shape), len(dst.shape))):\n    return False\n  # Identify input indices that are expanded vs. those that are preserved\n  # Expansion: input is 1, output is > 1.\n  # Preserved: input is > 1.\n  exp_indices, pre_indices = [], []\n  for i, dim in enumerate(src.shape):","sourceCodeStart":5286,"sourceCodeEnd":5322,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/fragmented_array.py#L5286-L5322","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","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."],"exampleFix":"// before\ncopy(src_smem_ref, dst_smem_ref, swizzle)  # both SMEM -> NotImplementedError\n// after\nregs = FragmentedArray.load_tiled(src_smem_ref, swizzle)\nregs.store_tiled(dst_smem_ref, swizzle)","handlingStrategy":"validation","validationCode":"from jax.experimental.mosaic.gpu import utils\n\ndef is_supported_copy(src, dst) -> bool:\n    src_smem, dst_smem = utils.is_smem_ref(src.type), utils.is_smem_ref(dst.type)\n    if src_smem == dst_smem:\n        return False  # only SMEM<->GMEM pairs are lowered\n    smem_ty, gmem_ty = (src.type, dst.type) if src_smem else (dst.type, src.type)\n    return smem_ty.rank == gmem_ty.rank + 2","typeGuard":null,"tryCatchPattern":"try:\n    copy(src, dst, swizzle)\nexcept NotImplementedError:\n    regs = FragmentedArray.load_tiled(src, swizzle)\n    regs.store_untiled(dst, optimized=False)","preventionTips":["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."],"tags":["jax","mosaic-gpu","gpu-kernel","shared-memory","not-implemented"],"backgroundTag":"unsupported-operation-not-implemented","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}