{"record":{"id":"c1e8cfcc18b9c421","repo":"jax-ml/jax","slug":"inline-mgpu-in-a-single-warp-context-only-supports-c1e8cf","errorCode":null,"errorMessage":"inline_mgpu in a single-warp context only supports scalar return types. Got shape={r.shape}.","messagePattern":"inline_mgpu in a single-warp context only supports scalar return types\\. Got shape=(.+?)\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/primitives.py","lineNumber":3651,"sourceCode":"@lowering.register_lowering_rule(inline_mgpu_p, mgpu.LoweringSemantics.Lane)\n@lowering.register_lowering_rule(inline_mgpu_p, *gpu_core.LANExWARP_SEMANTICS)\ndef _inline_mgpu_lowering_rule(\n    ctx: lowering.LoweringRuleContext,\n    *flat_args_and_transforms,\n    mgpu_fn: Callable[..., Any],\n    flat_arg_types,\n    flat_ret_ty,\n    pytree_args,\n    pytree_ref_transforms,\n    pytree_ret_ty,\n):\n  is_warp_semantics = (\n      ctx.module_ctx.primitive_semantics == gpu_core.PrimitiveSemantics.Warp\n  )\n  if is_warp_semantics:\n    for r in flat_ret_ty:\n      if isinstance(r, ShapeDtypeStruct) and r.shape:\n        raise ValueError(\n            \"inline_mgpu in a single-warp context only supports scalar return\"\n            f\" types. Got shape={r.shape}.\"\n        )\n\n  flat_transformed = _inline_mgpu_flat_transformed_args(\n      ctx,\n      flat_args_and_transforms,\n      flat_arg_types,\n      pytree_args,\n      pytree_ref_transforms,\n  )\n  args = jax.tree.unflatten(pytree_args, flat_transformed)\n  ret = mgpu_fn(ctx.launch_ctx, *args)\n  ret_leaves, ret_tree = jax.tree.flatten(\n      ret, lambda x: isinstance(x, mgpu.FragmentedArray)\n  )\n\n  if ret_tree != pytree_ret_ty:","sourceCodeStart":3633,"sourceCodeEnd":3669,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/primitives.py#L3633-L3669","documentation":"inline_mgpu executed under warp-level (single-warp) PrimitiveSemantics can only return scalars. Any ShapeDtypeStruct return value with a non-empty shape triggers this ValueError at abstract-evaluation time.","triggerScenarios":"Defining a kernel whose primitive semantics are Warp (e.g. called from a warp-specialized region) whose inline_mgpu callback returns an array (shape != ()).","commonSituations":"Porting a warp-group kernel to warp semantics without changing return handling; returning fragmented arrays from compute-in-warp callbacks.","solutions":["Change the callback to return only scalars (reduce to a scalar) or restructure so arrays are written to a ref","Switch the enclosing context to WarpGroup semantics if array returns are required","Return via an out_ref argument instead of the return value"],"exampleFix":"// before\ninline_mgpu(lambda x: x + 1, arr)  # returns array under Warp semantics\n// after\ninline_mgpu(lambda x: jnp.sum(x + 1), arr)  # scalar return","handlingStrategy":"validation","validationCode":"assert all(not getattr(r, 'shape', ()) for r in flat_ret_ty), 'warp context requires scalar returns'","typeGuard":"def is_scalar_ret(r): return not getattr(r, 'shape', ())","tryCatchPattern":null,"preventionTips":["Keep inline_mgpu callbacks scalar-only when called from warp contexts"],"tags":["jax","pallas","warp-semantics","scalar-return"],"backgroundTag":"unsupported-return-type","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}