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.

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

• Non-indexing transforms on GMEM refs are not implemented.
• Unsupported transform
• Commuting a `UntilingTransform` with a `ReshapeTransform`…
• Commuting a `UntilingTransform` with a `ReshapeTransform`…
• does not support and

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

Appendix: 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)