jax-ml/jax · error · ValueError

Transform mismatch: got {user_transforms}, expected {ty_tran

Error message

Transform mismatch: got {user_transforms}, expected {ty_transforms}

What it means

Raised when the transforms explicitly supplied by the user for an inline_mgpu argument do not match the transforms derived from the argument's abstract value (aval) via pallas_core.undo_transforms. This is a consistency check: the declared transform chain and the actual aval's baked-in transforms diverge.

Source

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

        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


@lowering.register_lowering_rule(inline_mgpu_p, mgpu.LoweringSemantics.Lane)
@lowering.register_lowering_rule(inline_mgpu_p, *gpu_core.LANExWARP_SEMANTICS)
def _inline_mgpu_lowering_rule(
    ctx: lowering.LoweringRuleContext,
    *flat_args_and_transforms,
    mgpu_fn: Callable[..., Any],
    flat_arg_types,
    flat_ret_ty,
    pytree_args,
    pytree_ref_transforms,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Pass the exact transforms reported in 'expected' (ty_transforms) instead of hand-built ones
Rebuild the ref from scratch with public APIs so transforms are attached automatically
Clear JIT caches and retry to rule out stale avals
Align JAX/pallas versions between components generating and consuming the transforms

Example fix

// before
inline_mgpu(fn, ref, transforms=(Transpose(0,1),))
// after
inline_mgpu(fn, ref, transforms=tuple(pallas_core.undo_transforms(aval, ref.transforms)))

Defensive patterns

Strategy: validation

Validate before calling

expected = tuple(pallas_core.undo_transforms(aval, t.transforms))
assert tuple(user_transforms) == expected, (user_transforms, expected)

Prevention

Never hand-construct transform tuples; derive them from the ref
Clear caches after changing kernel signatures

When it happens

Trigger: Calling inline_mgpu with user_transforms that differ in content or order from tuple(pallas_core.undo_transforms(aval, t.transforms)); typically from manually constructing transformed refs or stale cached avals.

Common situations: Mixing transform application order; constructing refs through internal APIs rather than public helpers; JAX version change altering transform normalization order.

Related errors

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