jax-ml/jax · error · ValueError

inline_mgpu_p only supports plgpu.ShapeDtypeStruct return ty

Error message

inline_mgpu_p only supports plgpu.ShapeDtypeStruct return types.

What it means

The @inline_mgpu decorator (used to embed raw Mosaic GPU / MLIR calls in Pallas kernels) requires return_type to be a pytree of ShapeDtypeStruct leaves (or empty). Any other return type annotation fails immediately.

Source

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

        x = mgpu.FragmentedArray.load_tiled(smem_ref)
        y = mgpu.FragmentedArray.splat(
            mgpu.c(1, x.mlir_dtype), shape=x.shape, layout=x.layout
        )
        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):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Wrap each returned value in plgpu.ShapeDtypeStruct(shape, dtype)
Pass return_type=() if the function returns nothing

Example fix

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

Defensive patterns

Strategy: validation

Validate before calling

import jax.tree_util as jtu
assert all(r is None or isinstance(r, ShapeDtypeStruct) for r in jtu.tree_leaves(return_type))

Type guard

def valid_return_type(rt):
    return all(isinstance(r, ShapeDtypeStruct) for r in jtu.tree_leaves(rt) if r is not None)

Prevention

Use ShapeDtypeStruct for every declared return leaf
Use () for no returns

When it happens

Trigger: Calling inline_mgpu(f, arg_types=..., return_type=int) or passing plain dtypes/tuples of dtypes instead of plgpu.ShapeDtypeStruct instances.

Common situations: Wrapping a helper that returns a python scalar or an mgpu type directly; annotating return_type with jax.numpy dtypes.

Related errors

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