jax-ml/jax · error · NotImplementedError

Not all transforms could be handled. Remaining transforms: {

Error message

Not all transforms could be handled. Remaining transforms: {user_transforms}.

What it means

Raised by JAX Pallas Mosaic GPU when an inline_mgpu call receives MemoryRefTransforms that the lowering infrastructure could not discharge. Under warp-group (WG) semantics every user-supplied transform must be explicitly handled; if any remain after processing, this NotImplementedError fires. It indicates the combination of transforms on the mgpu reference argument is unsupported.

Source

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

    if not isinstance(t, RefType):
      flat_transformed.append(a)
      assert transforms is None
      continue
    assert isinstance(aval, state.AbstractRef)
    assert isinstance(a, ir.Value)
    a, aval, user_transforms = lowering._handle_transforms(
        ctx,
        aval,
        a,
        transform_avals,
        transforms,
        handle_transposes=is_wg_semantics,
        allow_peer_refs=True,
    )

    if is_wg_semantics:
      if user_transforms:
        raise NotImplementedError(
            "Not all transforms could be handled. Remaining transforms:"
            f" {user_transforms}."
        )
    else:
      # Transforms that do not originate from a MemoryRefTransform are
      # applied implicitly (eg by emit-pipeline) and therefore we do not
      # expect the user to pass them to the type. The transforms not
      # passed by the user here will be discharged.
      ty_transforms = tuple(pallas_core.undo_transforms(aval, t.transforms))
      if ty_transforms != tuple(user_transforms):
        raise ValueError(
            f"Transform mismatch: got {user_transforms}, expected"
            f" {ty_transforms}"
        )
    flat_transformed.append(a)

  return flat_transformed

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove transforms (e.g. .T, reshape) from the ref passed to inline_mgpu and apply them manually inside the callback
2. Check that the transform types you use are handled by _handle_transforms (transposes under WG semantics)
3. Upgrade or downgrade JAX to a version whose mosaic_gpu primitives support your transform combination
4. File an issue with the remaining transforms list if the transform is a legitimate WG-supported case

Example fix

// before
inline_mgpu(fn, mgpu_ref.T)
// after
inline_mgpu(fn, mgpu_ref)  # transpose handled inside fn manually

assert not user_transforms or all_handled(user_transforms), f'unhandled transforms: {user_transforms}'

When it happens

Trigger: Calling pallas' inline_mgpu with a reference that carries transforms (transpose/broadcast/reshape via MemRefTransform) while the primitive is compiled with warp-group semantics (handle_transposes=True path), leaving unconsumed transforms in user_transforms.

Common situations: Passing a transposed or reshaped TMEM/SMEM ref into an inline_mgpu kernel; using newer Pallas APIs that attach transforms implicitly; version mismatches between jax and pallas where transform handling changed.

Related errors

• Unsupported transform: {type(transform)}
• Non-indexing transforms on GMEM refs are not implemented.
• 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/5d3bd9ba2dd31b53. Report an issue: GitHub.