jax-ml/jax · error · ValueError

inline_mgpu_p only supports only SomeLayout and RefType arg

Error message

inline_mgpu_p only supports only SomeLayout and RefType arg types.

What it means

inline_mgpu validates that every flattened entry of arg_types is either a layout specification (SomeLayout) or a RefType; anything else (plain dtypes, shapes) is rejected because the wrapper cannot map it to an MLIR argument.

Source

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

        return x + y

  Args:
    arg_types: A sequence of pytrees where the leaves are
      :class:`~jax.experimental.pallas.mosaic_gpu.RefType`\s or
      :class:`~jax.experimental.pallas.mosaic_gpu.Layout`\s for reference or
      array arguments respectively.
    return_type: A pytree where the leaves are
      :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):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Convert each argument type to plgpu.RefType(dtype, shape, layout) or a layout type
2. Keep ShapeDtypeStruct only for return_type

Example fix

# before
inline_mgpu(f, arg_types=[jnp.float32], return_type=...)
# after
inline_mgpu(f, arg_types=[plgpu.RefType(jnp.float32, (), plgpu.Layout())], return_type=...)

assert all(isinstance(t, (SomeLayout, RefType)) for t in jax.tree.flatten(arg_types)[0])

def valid_arg_types(at):
    return all(isinstance(t, (SomeLayout, RefType)) for t in jax.tree.flatten(at)[0])

When it happens

Trigger: Passing arg_types entries like jnp.float32 or ShapeDtypeStruct (which are only valid for return_type) to inline_mgpu.

Common situations: Copying return_type conventions into arg_types; mixing up which side uses ShapeDtypeStruct vs RefType/layout types.

Related errors

• inline_mgpu_p only supports plgpu.ShapeDtypeStruct return ty
• Unexpected type {ty} for value {v}
• MemoryRef type must be a ShapedArray, got {type(self.inner_a
• Unsupported block dimension type: {type(dim)}. Allowed types
• Unsupported block dimension type: {type(bd)}

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