{"record":{"id":"a469ef39ff208d8f","repo":"jax-ml/jax","slug":"transforms-are-not-yet-implemented-for-warpgroup-s","errorCode":null,"errorMessage":"Transforms are not yet implemented for warpgroup semantics","messagePattern":"Transforms are not yet implemented for warpgroup semantics","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/lowering.py","lineNumber":2239,"sourceCode":"def _get_lowering_rule_wg(\n    ctx: LoweringRuleContext, x_ref, *leaves, tree, optimized=True\n):\n  if not isinstance(x_ref, ir.Value) and isinstance(x_ref, ir.MemRefType):\n    raise TypeError(f\"Can only load from references (got {x_ref}).\")\n  shape = ctx.avals_out[0].shape\n  if shape and ctx.module_ctx.primitive_semantics == gpu_core.PrimitiveSemantics.Warp:\n    raise ValueError(\"Can only load scalars in warp-level code.\")\n\n  transforms = jax.tree.unflatten(tree, leaves)\n  assert isinstance(ctx.avals_in[0], state_types.AbstractRef)\n  transform_avals = jax.tree.unflatten(tree, ctx.avals_in[1:])\n  x_ref, _, transforms = _handle_transforms(\n      ctx, ctx.avals_in[0], x_ref, transform_avals, transforms,\n      allow_peer_refs=True\n  )\n\n  if transforms:\n    raise NotImplementedError(\n        \"Transforms are not yet implemented for warpgroup semantics\"\n    )\n\n  assert isinstance(x_ref, ir.Value)\n  shape = ctx.avals_out[0].shape\n  if shape:\n    return mgpu.dialect.vector_load(x_ref, optimized=optimized)\n  else:\n    return memref_dialect.load(x_ref, [])\n\n\n@register_lowering_rule(sp.swap_p, mgpu.LoweringSemantics.Lane)\n@register_lowering_rule(sp.swap_p, *gpu_core.LANExWARP_SEMANTICS)\ndef _swap_lowering_rule(\n    ctx: LoweringRuleContext, x_ref, value, *leaves, tree\n):\n  if isinstance(x_ref, tcgen05.TMEMRef):\n    raise RuntimeError(","sourceCodeStart":2221,"sourceCodeEnd":2257,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/lowering.py#L2221-L2257","documentation":"The warpgroup-semantics load lowering rule rejects any reference transforms (slicing/tiling/swizzle descriptors) on the loaded reference. Warpgroup lowering only supports plain shared/global memory loads without transform chains.","triggerScenarios":"Passing a transformed ref (e.g. x_ref[2:5, :] or a swizzled/tilted ref via plgpu transform descriptors) into a load lowered with LoweringSemantics.Warpgroup.","commonSituations":"Reusing lane-level kernel code that indexes into refs inside a warpgroup-scoped region (e.g. around tcgen05 MMA blocks); dynamic_slice or slicing on a ref under warpgroup semantics.","solutions":["Apply the slicing to the underlying memory before entering warpgroup semantics, or use a separately allocated ref","Express the operation with explicit TMEM/shared-memory loads (e.g. mgpu/load from tmem) that don't require transforms","Restructure so the value is loaded without any transform chain"],"exampleFix":"# before\nwith plgpu.warp_group_semantics():\n  v = x_ref[0:128, :]  # transform under warpgroup\n# after\nv = x_ref[0:128, :]  # slice at lane semantics / before entering wg region","handlingStrategy":"validation","validationCode":null,"typeGuard":null,"tryCatchPattern":null,"preventionTips":["Never slice/index refs inside warpgroup regions","Load plain (untransformed) refs in wg code; do transforms outside"],"tags":["jax","pallas","mosaic-gpu","transforms","warpgroup"],"backgroundTag":"unsupported-operation-language-feature","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}