jax-ml/jax · error · ValueError

inline_mgpu in a single-warp context only supports scalar ar

Error message

inline_mgpu in a single-warp context only supports scalar arrays (and Refs). Got {aval}.

What it means

With single-warp (lane) semantics, each GPU lane executes the inline_mgpu function, so non-ref arguments must be scalars replicated across lanes. A ShapedArray with a non-empty shape passed where a non-RefType was declared triggers this error.

Source

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

  )

  if is_wg_semantics:
    flat_args = [
        lowering._ensure_ir_value(a, aval.dtype) if not isinstance(t, RefType) else a
        for a, aval, t in zip(flat_args, flat_arg_avals, flat_arg_types)
    ]
  else:
    flat_args = [
        lowering._ensure_fa(a, aval.dtype) if not isinstance(t, RefType) else a
        for a, aval, t in zip(flat_args, flat_arg_avals, flat_arg_types)
    ]

  for a, aval, t in zip(flat_args, flat_arg_avals, flat_arg_types):
    if not is_wg_semantics:
      _type_check_mgpu_lane_semantics(a, t)
    if is_warp_semantics and not isinstance(t, RefType):
      if not isinstance(aval, jax_core.ShapedArray) or aval.shape:
        raise ValueError(
            "inline_mgpu in a single-warp context only supports scalar"
            f" arrays (and Refs). Got {aval}."
        )

  flat_transformed : list[ir.Value | mgpu.FragmentedArray] = []
  for a, aval, t, transforms, transform_avals in zip(
      flat_args,
      flat_arg_avals,
      flat_arg_types,
      ref_transforms,
      ref_transform_avals,
      strict=True,
  ):
    if not isinstance(t, RefType):
      flat_transformed.append(a)
      assert transforms is None
      continue
    assert isinstance(aval, state.AbstractRef)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass array data through RefType (SMEM) arguments instead of by-value arrays
2. Keep by-value arguments as scalars (ShapeDtypeStruct with shape ())
3. Use warp-group semantics if arrays must be passed by value

Example fix

# before
inline_mgpu(f, arg_types=[Layout(...)], semantics=LANE)
# after
inline_mgpu(f, arg_types=[RefType(dt, shape, layout)], semantics=LANE)

assert isinstance(t, RefType) or (isinstance(aval, jax_core.ShapedArray) and not aval.shape)

def warp_safe_arg(t, aval):
    return isinstance(t, RefType) or (isinstance(aval, jax_core.ShapedArray) and aval.shape == ())

When it happens

Trigger: Using inline_mgpu under warp semantics and passing a tensor-shaped argument declared as a layout type instead of a scalar ShapeDtypeStruct/RefType.

Common situations: Porting warp-group kernels to warp semantics; accidentally declaring arrays (not refs) as inputs in warp mode.

Related errors

• trace_value requires a scalar value, got shape {value.shape}
• Can only store scalars in warp-level lowering.
• inline_mgpu_p only supports plgpu.ShapeDtypeStruct return ty
• inline_mgpu_p only supports only SomeLayout and RefType arg
• Mismatched type shape: {treedef} != {treedef_ty}

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