jax-ml/jax · error · ValueError

Mismatched type shape: {treedef} != {treedef_ty}

Error message

Mismatched type shape: {treedef} != {treedef_ty}

What it means

The function wrapped by inline_mgpu must be called with arguments whose pytree structure exactly matches the arg_types structure passed at decoration time. The wrapper flattens both and compares treedefs; any difference (extra/missing args, different nesting) raises this error.

Source

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

      :class:`~jax.experimental.pallas.mosaic_gpu.ShapeDtypeStruct`\s
      representing the arrays returned by the decorated function.
  """
  flat_arg_types, treedef_ty = jax.tree.flatten(tuple(arg_types))
  flat_ret_ty, pytree_ret_ty = jax.tree.flatten(return_type)
  if return_type and not all(isinstance(r, ShapeDtypeStruct) for r in flat_ret_ty):
    raise ValueError(
        "inline_mgpu_p only supports plgpu.ShapeDtypeStruct return types."
    )
  if not all(isinstance(r, (SomeLayout, RefType)) for r in flat_arg_types):
    raise ValueError(
        "inline_mgpu_p only supports only SomeLayout and RefType arg types."
    )

  def inner(f):
    def wrapper(*args):
      flat_args, treedef = jax.tree.flatten(tuple(args))
      if treedef != treedef_ty:
        raise ValueError(f"Mismatched type shape: {treedef} != {treedef_ty}")

      # Strip the transforms from the refs since they will be recorded in
      # the types.
      ref_transforms: list[Any] = []
      raw_flat_args = []
      for a, t in zip(flat_args, flat_arg_types):
        if isinstance(a, state_types.TransformedRef) and isinstance(t, RefType):
          raw_flat_args.append(a.ref)
          ref_transforms.append(a.transforms)
        elif isinstance(aval := jax_core.typeof(a), jax_core.ShapedArray) and isinstance(t, SomeLayout):
          raw_flat_args.append(a)
          ref_transforms.append(None)
        elif isinstance(aval, state.AbstractRef) and isinstance(t, RefType):
          raw_flat_args.append(a)
          ref_transforms.append(())
        else:
          raise ValueError(f"Mismatched type: {a, t}")

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Call the wrapped function with exactly the same arg structure as arg_types
  2. Update arg_types to match the new call signature

Example fix

# before
f = inline_mgpu(impl, arg_types=[ty_a, ty_b], return_type=...)
f(a)
# after
f(a, b)
Defensive patterns

Strategy: validation

Validate before calling

assert jax.tree.structure(args) == jax.tree.structure(arg_types_tuple)

Prevention

When it happens

Trigger: Decorating with arg_types=[RefType, RefType] but calling f with one argument, or passing args nested differently ((a,b),c vs a,b,c).

Common situations: Refactoring the wrapped function signature after decoration; passing optional args conditionally.

Related errors

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