{"record":{"id":"c21f54cf81f6940e","repo":"jax-ml/jax","slug":"inline-mgpu-p-only-supports-plgpu-shapedtypestruct","errorCode":null,"errorMessage":"inline_mgpu_p only supports plgpu.ShapeDtypeStruct return types.","messagePattern":"inline_mgpu_p only supports plgpu\\.ShapeDtypeStruct return types\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/primitives.py","lineNumber":3441,"sourceCode":"        x = mgpu.FragmentedArray.load_tiled(smem_ref)\n        y = mgpu.FragmentedArray.splat(\n            mgpu.c(1, x.mlir_dtype), shape=x.shape, layout=x.layout\n        )\n        return x + y\n\n  Args:\n    arg_types: A sequence of pytrees where the leaves are\n      :class:`~jax.experimental.pallas.mosaic_gpu.RefType`\\s or\n      :class:`~jax.experimental.pallas.mosaic_gpu.Layout`\\s for reference or\n      array arguments respectively.\n    return_type: A pytree where the leaves are\n      :class:`~jax.experimental.pallas.mosaic_gpu.ShapeDtypeStruct`\\s\n      representing the arrays returned by the decorated function.\n  \"\"\"\n  flat_arg_types, treedef_ty = jax.tree.flatten(tuple(arg_types))\n  flat_ret_ty, pytree_ret_ty = jax.tree.flatten(return_type)\n  if return_type and not all(isinstance(r, ShapeDtypeStruct) for r in flat_ret_ty):\n    raise ValueError(\n        \"inline_mgpu_p only supports plgpu.ShapeDtypeStruct return types.\"\n    )\n  if not all(isinstance(r, (SomeLayout, RefType)) for r in flat_arg_types):\n    raise ValueError(\n        \"inline_mgpu_p only supports only SomeLayout and RefType arg types.\"\n    )\n\n  def inner(f):\n    def wrapper(*args):\n      flat_args, treedef = jax.tree.flatten(tuple(args))\n      if treedef != treedef_ty:\n        raise ValueError(f\"Mismatched type shape: {treedef} != {treedef_ty}\")\n\n      # Strip the transforms from the refs since they will be recorded in\n      # the types.\n      ref_transforms: list[Any] = []\n      raw_flat_args = []\n      for a, t in zip(flat_args, flat_arg_types):","sourceCodeStart":3423,"sourceCodeEnd":3459,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/primitives.py#L3423-L3459","documentation":"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.","triggerScenarios":"Calling inline_mgpu(f, arg_types=..., return_type=int) or passing plain dtypes/tuples of dtypes instead of plgpu.ShapeDtypeStruct instances.","commonSituations":"Wrapping a helper that returns a python scalar or an mgpu type directly; annotating return_type with jax.numpy dtypes.","solutions":["Wrap each returned value in plgpu.ShapeDtypeStruct(shape, dtype)","Pass return_type=() if the function returns nothing"],"exampleFix":"# before\ninline_mgpu(f, arg_types=[...], return_type=jnp.float32)\n# after\ninline_mgpu(f, arg_types=[...], return_type=plgpu.ShapeDtypeStruct((), jnp.float32))","handlingStrategy":"validation","validationCode":"import jax.tree_util as jtu\nassert all(r is None or isinstance(r, ShapeDtypeStruct) for r in jtu.tree_leaves(return_type))","typeGuard":"def valid_return_type(rt):\n    return all(isinstance(r, ShapeDtypeStruct) for r in jtu.tree_leaves(rt) if r is not None)","tryCatchPattern":null,"preventionTips":["Use ShapeDtypeStruct for every declared return leaf","Use () for no returns"],"tags":["jax","pallas","mosaic-gpu","inline-mgpu","type-validation"],"backgroundTag":"invalid-type-annotation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}