{"record":{"id":"036520c179c259bd","repo":"jax-ml/jax","slug":"inline-mgpu-p-only-supports-only-somelayout-and-re","errorCode":null,"errorMessage":"inline_mgpu_p only supports only SomeLayout and RefType arg types.","messagePattern":"inline_mgpu_p only supports only SomeLayout and RefType arg types\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/primitives.py","lineNumber":3445,"sourceCode":"        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):\n        if isinstance(a, state_types.TransformedRef) and isinstance(t, RefType):\n          raw_flat_args.append(a.ref)\n          ref_transforms.append(a.transforms)\n        elif isinstance(aval := jax_core.typeof(a), jax_core.ShapedArray) and isinstance(t, SomeLayout):","sourceCodeStart":3427,"sourceCodeEnd":3463,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/primitives.py#L3427-L3463","documentation":"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.","triggerScenarios":"Passing arg_types entries like jnp.float32 or ShapeDtypeStruct (which are only valid for return_type) to inline_mgpu.","commonSituations":"Copying return_type conventions into arg_types; mixing up which side uses ShapeDtypeStruct vs RefType/layout types.","solutions":["Convert each argument type to plgpu.RefType(dtype, shape, layout) or a layout type","Keep ShapeDtypeStruct only for return_type"],"exampleFix":"# before\ninline_mgpu(f, arg_types=[jnp.float32], return_type=...)\n# after\ninline_mgpu(f, arg_types=[plgpu.RefType(jnp.float32, (), plgpu.Layout())], return_type=...)","handlingStrategy":"type-guard","validationCode":"assert all(isinstance(t, (SomeLayout, RefType)) for t in jax.tree.flatten(arg_types)[0])","typeGuard":"def valid_arg_types(at):\n    return all(isinstance(t, (SomeLayout, RefType)) for t in jax.tree.flatten(at)[0])","tryCatchPattern":null,"preventionTips":["Remember: args use RefType/layout, returns use ShapeDtypeStruct"],"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"}