jax-ml/jax · error · ValueError

inline_mgpu in a single-warp context only supports scalar re

Error message

inline_mgpu in a single-warp context only supports scalar return types. Got shape={r.shape}.

What it means

inline_mgpu executed under warp-level (single-warp) PrimitiveSemantics can only return scalars. Any ShapeDtypeStruct return value with a non-empty shape triggers this ValueError at abstract-evaluation time.

Source

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

@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,
    pytree_ret_ty,
):
  is_warp_semantics = (
      ctx.module_ctx.primitive_semantics == gpu_core.PrimitiveSemantics.Warp
  )
  if is_warp_semantics:
    for r in flat_ret_ty:
      if isinstance(r, ShapeDtypeStruct) and r.shape:
        raise ValueError(
            "inline_mgpu in a single-warp context only supports scalar return"
            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:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Change the callback to return only scalars (reduce to a scalar) or restructure so arrays are written to a ref
2. Switch the enclosing context to WarpGroup semantics if array returns are required
3. Return via an out_ref argument instead of the return value

Example fix

// before
inline_mgpu(lambda x: x + 1, arr)  # returns array under Warp semantics
// after
inline_mgpu(lambda x: jnp.sum(x + 1), arr)  # scalar return

assert all(not getattr(r, 'shape', ()) for r in flat_ret_ty), 'warp context requires scalar returns'

def is_scalar_ret(r): return not getattr(r, 'shape', ())

When it happens

Trigger: Defining a kernel whose primitive semantics are Warp (e.g. called from a warp-specialized region) whose inline_mgpu callback returns an array (shape != ()).

Common situations: Porting a warp-group kernel to warp semantics without changing return handling; returning fragmented arrays from compute-in-warp callbacks.

Related errors

• Can only store scalars in warp-level lowering.
• inline_mgpu in a single-warp context only supports scalar ar
• This gmm kernel only supports either (m, k) x (g, k, n) -> (
• Group sizes {group_sizes.shape=} must match first dimension
• Explicit sharding is not currently supported in the pallas-t

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