{"record":{"id":"998251a8c526f489","repo":"jax-ml/jax","slug":"unexpected-unhandled-transforms-remaining-ref-tr","errorCode":null,"errorMessage":"Unexpected unhandled transforms: {remaining_ref_transforms}","messagePattern":"Unexpected unhandled transforms: (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/primitives.py","lineNumber":653,"sourceCode":"      ref_transforms,\n      handle_transposes=True,\n  )\n\n  base_index = _get_barrier_base_index(barrier_ref_aval, barrier_transforms)\n  if base_index is not None:\n    barrier = barrier[base_index]\n\n  cluster_idx_val = lowering._as_index(cluster_idx)\n  gpu_cluster_dim = lowering._resolve_cluster_axis(ctx.module_ctx.axis_names, cluster_dim)\n\n  shape = ctx.avals_in[0].shape\n  dtype = ctx.avals_in[0].dtype\n  if not shape:\n    raise NotImplementedError(\"Scalars are not supported in async_store_smem\")\n\n  if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Warpgroup:\n    if remaining_ref_transforms:\n      raise ValueError(f\"Unexpected unhandled transforms: {remaining_ref_transforms}\")\n    assert isinstance(barrier, mgpu.DialectBarrierRef)\n    cluster_idx_i32 = arith_dialect.index_cast(\n        ir.IntegerType.get_signless(32), cluster_idx_val\n    )\n    atomic_type = None\n    if atomic is not None:\n      atomic_type = _atomic_op_type_to_int(AtomicOpType(atomic))\n    mgpu.dialect.async_store_smem(\n        src,\n        ref_smem,\n        barrier.as_barrier_memref(),\n        gpu_cluster_dim.value,\n        cluster_idx_i32,\n        atomic_type=atomic_type,\n        optimized=optimized,\n    )\n    return ()\n","sourceCodeStart":635,"sourceCodeEnd":671,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/primitives.py#L635-L671","documentation":"In the Warpgroup lowering path of _async_store_smem_lowering, any ref transforms that remain unprocessed at that point are unexpected and raise ValueError. This is an internal invariant: by this stage all transforms should have been consumed (validated as unswizzle+untiling later in Warp semantics).","triggerScenarios":"Passing transform combinations on the SMEM ref that the warpgroup path doesn't consume — e.g. extra slicing transforms on top of the required unswizzle/tiling, or transforms in an unexpected order.","commonSituations":"Using experimental transform APIs on refs; hitting an internal bug after upgrading JAX where a new transform type isn't handled yet; combining swizzle with additional indexing on async-store targets.","solutions":["Reduce the ref to only the standard tiled+swizzled form (no extra transforms) before async_store_smem","Apply slicing to the value instead of the ref","If the transforms look legitimate, file a JAX issue — this is an internal invariant violation"],"exampleFix":"# before\nasync_store_smem(smem[block_indices], x, barrier)  # extra slicing\n# after\nasync_store_smem(smem, x_sliced, barrier)","handlingStrategy":"type-guard","validationCode":null,"typeGuard":"def only_unswizzle_untiling(transforms):\n  from jax._src.pallas import gpu_core\n  return (len(transforms) == 2 and\n          isinstance(transforms[0], gpu_core.UnswizzleRef) and\n          isinstance(transforms[1], gpu_core.UntilingTransform))","tryCatchPattern":"try:\n    async_store_smem(...)\nexcept ValueError as e:\n    if 'Unhandled transforms' in str(e):\n        smem[...] = value  # fallback synchronous store\n    else:\n        raise","preventionTips":["Avoid extra ref transforms beyond the standard tiled+swizzled form","Report internal invariant hits to the JAX repo with a repro"],"tags":["jax","pallas","mosaic-gpu","internal-invariant","transforms","async-store"],"backgroundTag":"invalid-argument-combination","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}