jax-ml/jax · error · ValueError

inline_mgpu_p return type tree mismatch: {ret} != {return_ty

Error message

inline_mgpu_p return type tree mismatch: {ret} != {return_type}

What it means

The pytree structure of values returned by the user's inline_mgpu function must exactly match the declared return type tree. Flattening the actual return with FragmentedArray leaves yields a different tree structure than pytree_ret_ty, so this ValueError is raised.

Source

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

            f" types. Got shape={r.shape}."
        )

  flat_transformed = _inline_mgpu_flat_transformed_args(
      ctx,
      flat_args_and_transforms,
      flat_arg_types,
      pytree_args,
      pytree_ref_transforms,
  )
  args = jax.tree.unflatten(pytree_args, flat_transformed)
  ret = mgpu_fn(ctx.launch_ctx, *args)
  ret_leaves, ret_tree = jax.tree.flatten(
      ret, lambda x: isinstance(x, mgpu.FragmentedArray)
  )

  if ret_tree != pytree_ret_ty:
    return_type = jax.tree.unflatten(pytree_ret_ty, flat_ret_ty)
    raise ValueError(
        f"inline_mgpu_p return type tree mismatch: {ret} != {return_type}"
    )

  for ty, r in zip(flat_ret_ty, ret_leaves):
    _type_check_mgpu_lane_semantics(r, ty)

  return ret_leaves


def _ref_type_to_transforms(ref_type: RefType) -> ir.ArrayAttr:
  """Returns the Mosaic GPU transforms for the given ref type."""
  transform_attrs = [gpu_core.to_transform_attr(t)
                     for t in ref_type.transforms]
  return ir.ArrayAttr.get(transform_attrs)


def _custom_primitive_in_specs(
    ctx: lowering.LoweringRuleContext,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Make the callback return exactly the declared pytree structure (same nesting and leaf count)
Update the declared return-type tree whenever the callback's return expression changes
Avoid conditional returns with different structures; return a fixed tuple shape

Defensive patterns

Strategy: type-guard

Validate before calling

import jax
assert jax.tree.structure(fn(*args)).num_leaves == len(flat_ret_ty)

Type guard

def ret_tree_matches(fn, args, ret_tree):
    return jax.tree.structure(fn(*args)) == ret_tree

Prevention

Return a fixed-structure tuple from callbacks
Update declared out types whenever the callback return changes

When it happens

Trigger: The callback returns a differently nested structure (e.g. a tuple vs a single array, wrong leaf count) than the avals passed as flat_ret_ty / their pytree structure.

Common situations: Returning (array,) vs array; changing the callback signature or return shape without updating the declared output avals; conditional returns producing inconsistent trees.

Related errors

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