jax-ml/jax · error · ValueError

Unexpected unhandled transforms: {remaining_ref_transforms}

Error message

Unexpected unhandled transforms: {remaining_ref_transforms}

What it means

In the Warpgroup lowering path of _async_store_smem_lowering, any ref transforms that remain unprocessed at that point are unexpected and raise ValueError. This is an internal invariant: by this stage all transforms should have been consumed (validated as unswizzle+untiling later in Warp semantics).

Source

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

      ref_transforms,
      handle_transposes=True,
  )

  base_index = _get_barrier_base_index(barrier_ref_aval, barrier_transforms)
  if base_index is not None:
    barrier = barrier[base_index]

  cluster_idx_val = lowering._as_index(cluster_idx)
  gpu_cluster_dim = lowering._resolve_cluster_axis(ctx.module_ctx.axis_names, cluster_dim)

  shape = ctx.avals_in[0].shape
  dtype = ctx.avals_in[0].dtype
  if not shape:
    raise NotImplementedError("Scalars are not supported in async_store_smem")

  if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Warpgroup:
    if remaining_ref_transforms:
      raise ValueError(f"Unexpected unhandled transforms: {remaining_ref_transforms}")
    assert isinstance(barrier, mgpu.DialectBarrierRef)
    cluster_idx_i32 = arith_dialect.index_cast(
        ir.IntegerType.get_signless(32), cluster_idx_val
    )
    atomic_type = None
    if atomic is not None:
      atomic_type = _atomic_op_type_to_int(AtomicOpType(atomic))
    mgpu.dialect.async_store_smem(
        src,
        ref_smem,
        barrier.as_barrier_memref(),
        gpu_cluster_dim.value,
        cluster_idx_i32,
        atomic_type=atomic_type,
        optimized=optimized,
    )
    return ()

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Reduce the ref to only the standard tiled+swizzled form (no extra transforms) before async_store_smem
2. Apply slicing to the value instead of the ref
3. If the transforms look legitimate, file a JAX issue — this is an internal invariant violation

Example fix

# before
async_store_smem(smem[block_indices], x, barrier)  # extra slicing
# after
async_store_smem(smem, x_sliced, barrier)

def only_unswizzle_untiling(transforms):
  from jax._src.pallas import gpu_core
  return (len(transforms) == 2 and
          isinstance(transforms[0], gpu_core.UnswizzleRef) and
          isinstance(transforms[1], gpu_core.UntilingTransform))

try:
    async_store_smem(...)
except ValueError as e:
    if 'Unhandled transforms' in str(e):
        smem[...] = value  # fallback synchronous store
    else:
        raise

When it happens

Trigger: Passing transform combinations on the SMEM ref that the warpgroup path doesn't consume — e.g. extra slicing transforms on top of the required unswizzle/tiling, or transforms in an unexpected order.

Common situations: Using experimental transform APIs on refs; hitting an internal bug after upgrading JAX where a new transform type isn't handled yet; combining swizzle with additional indexing on async-store targets.

Related errors

• Transpose cannot be moved before a tiling transform when it
• Commuting a `UntilingTransform` with a `ReshapeTransform` is
• Commuting a `UntilingTransform` with a `ReshapeTransform` is
• Unsupported transform: {type(transform)}
• Transforms are not yet implemented for warpgroup semantics

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