{"record":{"id":"49a5098728098777","repo":"jax-ml/jax","slug":"inline-mgpu-in-a-single-warp-context-only-supports","errorCode":null,"errorMessage":"inline_mgpu in a single-warp context only supports scalar arrays (and Refs). Got {aval}.","messagePattern":"inline_mgpu in a single-warp context only supports scalar arrays \\(and Refs\\)\\. Got (.+?)\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/primitives.py","lineNumber":3581,"sourceCode":"  )\n\n  if is_wg_semantics:\n    flat_args = [\n        lowering._ensure_ir_value(a, aval.dtype) if not isinstance(t, RefType) else a\n        for a, aval, t in zip(flat_args, flat_arg_avals, flat_arg_types)\n    ]\n  else:\n    flat_args = [\n        lowering._ensure_fa(a, aval.dtype) if not isinstance(t, RefType) else a\n        for a, aval, t in zip(flat_args, flat_arg_avals, flat_arg_types)\n    ]\n\n  for a, aval, t in zip(flat_args, flat_arg_avals, flat_arg_types):\n    if not is_wg_semantics:\n      _type_check_mgpu_lane_semantics(a, t)\n    if is_warp_semantics and not isinstance(t, RefType):\n      if not isinstance(aval, jax_core.ShapedArray) or aval.shape:\n        raise ValueError(\n            \"inline_mgpu in a single-warp context only supports scalar\"\n            f\" arrays (and Refs). Got {aval}.\"\n        )\n\n  flat_transformed : list[ir.Value | mgpu.FragmentedArray] = []\n  for a, aval, t, transforms, transform_avals in zip(\n      flat_args,\n      flat_arg_avals,\n      flat_arg_types,\n      ref_transforms,\n      ref_transform_avals,\n      strict=True,\n  ):\n    if not isinstance(t, RefType):\n      flat_transformed.append(a)\n      assert transforms is None\n      continue\n    assert isinstance(aval, state.AbstractRef)","sourceCodeStart":3563,"sourceCodeEnd":3599,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/primitives.py#L3563-L3599","documentation":"With single-warp (lane) semantics, each GPU lane executes the inline_mgpu function, so non-ref arguments must be scalars replicated across lanes. A ShapedArray with a non-empty shape passed where a non-RefType was declared triggers this error.","triggerScenarios":"Using inline_mgpu under warp semantics and passing a tensor-shaped argument declared as a layout type instead of a scalar ShapeDtypeStruct/RefType.","commonSituations":"Porting warp-group kernels to warp semantics; accidentally declaring arrays (not refs) as inputs in warp mode.","solutions":["Pass array data through RefType (SMEM) arguments instead of by-value arrays","Keep by-value arguments as scalars (ShapeDtypeStruct with shape ())","Use warp-group semantics if arrays must be passed by value"],"exampleFix":"# before\ninline_mgpu(f, arg_types=[Layout(...)], semantics=LANE)\n# after\ninline_mgpu(f, arg_types=[RefType(dt, shape, layout)], semantics=LANE)","handlingStrategy":"type-guard","validationCode":"assert isinstance(t, RefType) or (isinstance(aval, jax_core.ShapedArray) and not aval.shape)","typeGuard":"def warp_safe_arg(t, aval):\n    return isinstance(t, RefType) or (isinstance(aval, jax_core.ShapedArray) and aval.shape == ())","tryCatchPattern":null,"preventionTips":["Route arrays through RefType in warp mode","Reserve non-scalar by-value args for warp-group semantics"],"tags":["jax","pallas","inline-mgpu","warp-semantics","scalar-required"],"backgroundTag":"rank-mismatch-in-kernel-argument","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}