jax-ml/jax · error · NotImplementedError

Non-indexing transforms on GMEM refs are not implemented.

Error message

Non-indexing transforms on GMEM refs are not implemented.

What it means

While walking transforms on a GMEM reference, _extract_gmem_copy_params only understands NDIndexer transforms (plus the multicast transform handled above). Any other transform type on the ref (slicing helpers, swizzle, permutation, etc.) reaches the else branch and raises NotImplementedError.

Source

Thrown at jax/_src/pallas/mosaic_gpu/primitives.py:460

            "Multicast refs are not supported by this primitive."
        )
      if (mesh_info := ctx.module_ctx.mesh_info) is None:
        raise ValueError(
            "JAX device mesh is required by multicast copies, but not defined."
            " Use jax.set_mesh."
        )
      if set(transform.collective_axes) != set(mesh_info.axis_names):
        raise NotImplementedError(
            "Only collective_axes that include all JAX device mesh  axes are"
            f" supported, but got {transform.collective_axes}. Make sure to"
            f" pass collective_axes={mesh_info.axis_names}"
        )
      peer_id = mgpu.GLOBAL_BROADCAST
      continue
    elif isinstance(transform, indexing.NDIndexer):
      indexers.append(transform)
    else:
      raise NotImplementedError(
          "Non-indexing transforms on GMEM refs are not implemented.")
  if indexers:
    indexer = lowering.merge_indexers(indexers)
    gmem_slice = lowering._ndindexer_indices(indexer, allow_arrays=True)
  else:
    gmem_slice = ()
  return dict(
      gmem_slice=gmem_slice,
      gmem_peer_id=peer_id,
  )


def _extract_smem_copy_params(aval, transforms):
  if not transforms:
    return {}
  # Split off swizzling, if present
  match transforms:
    case [gpu_core.UnswizzleRef(swizzle), *transforms]:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Index GMEM refs only with standard Pallas indexing (block_indices / NDIndexer-producing ops) before the copy
2. Materialize the transformed value into an intermediate SMEM buffer and copy that instead
3. Check the JAX version/changelog — support for additional transforms may have been added in newer releases

Example fix

# before
copy_gmem_to_smem(gref.unswizzle(swizzle), smem)
# after
smem[...] = gref[...]  # plain indexing
copy_gmem_to_smem(gref, smem)

from jax._src.pallas.mosaic_gpu import indexing
def refs_have_only_indexer_transforms(ref_transforms):
  return all(t is None or isinstance(t, indexing.NDIndexer) or _is_multicast(t)
             for t in ref_transforms)

When it happens

Trigger: Passing a GMEM BlockRef with a non-indexing transform (e.g. a sliced/swizzled/permuted ref produced by ref.swap_dims, unswizzle, or similar) to copy_gmem_to_smem/copy_smem_to_gmem/prefetch in a Mosaic GPU kernel.

Common situations: Trying to reuse SMEM-style layout transforms (tiled/swizzled refs) on GMEM refs; applying fancy indexing utilities not based on indexing.NDIndexer; upgrading JAX where transform kinds expanded beyond what the copy path handles.

Related errors

• Unsupported transform: {type(transform)}
• Not all transforms could be handled. Remaining transforms: {
• Unsupported dtype: {ref.dtype}
• Only SMEM and TMEM refs are supported.
• Transpose cannot be moved before a tiling transform when it

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