{"record":{"id":"5d3bd9ba2dd31b53","repo":"jax-ml/jax","slug":"not-all-transforms-could-be-handled-remaining-tra","errorCode":null,"errorMessage":"Not all transforms could be handled. Remaining transforms: {user_transforms}.","messagePattern":"Not all transforms could be handled\\. Remaining transforms: (.+?)\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/primitives.py","lineNumber":3613,"sourceCode":"    if not isinstance(t, RefType):\n      flat_transformed.append(a)\n      assert transforms is None\n      continue\n    assert isinstance(aval, state.AbstractRef)\n    assert isinstance(a, ir.Value)\n    a, aval, user_transforms = lowering._handle_transforms(\n        ctx,\n        aval,\n        a,\n        transform_avals,\n        transforms,\n        handle_transposes=is_wg_semantics,\n        allow_peer_refs=True,\n    )\n\n    if is_wg_semantics:\n      if user_transforms:\n        raise NotImplementedError(\n            \"Not all transforms could be handled. Remaining transforms:\"\n            f\" {user_transforms}.\"\n        )\n    else:\n      # Transforms that do not originate from a MemoryRefTransform are\n      # applied implicitly (eg by emit-pipeline) and therefore we do not\n      # expect the user to pass them to the type. The transforms not\n      # passed by the user here will be discharged.\n      ty_transforms = tuple(pallas_core.undo_transforms(aval, t.transforms))\n      if ty_transforms != tuple(user_transforms):\n        raise ValueError(\n            f\"Transform mismatch: got {user_transforms}, expected\"\n            f\" {ty_transforms}\"\n        )\n    flat_transformed.append(a)\n\n  return flat_transformed\n","sourceCodeStart":3595,"sourceCodeEnd":3631,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/primitives.py#L3595-L3631","documentation":"Raised by JAX Pallas Mosaic GPU when an inline_mgpu call receives MemoryRefTransforms that the lowering infrastructure could not discharge. Under warp-group (WG) semantics every user-supplied transform must be explicitly handled; if any remain after processing, this NotImplementedError fires. It indicates the combination of transforms on the mgpu reference argument is unsupported.","triggerScenarios":"Calling pallas' inline_mgpu with a reference that carries transforms (transpose/broadcast/reshape via MemRefTransform) while the primitive is compiled with warp-group semantics (handle_transposes=True path), leaving unconsumed transforms in user_transforms.","commonSituations":"Passing a transposed or reshaped TMEM/SMEM ref into an inline_mgpu kernel; using newer Pallas APIs that attach transforms implicitly; version mismatches between jax and pallas where transform handling changed.","solutions":["Remove transforms (e.g. .T, reshape) from the ref passed to inline_mgpu and apply them manually inside the callback","Check that the transform types you use are handled by _handle_transforms (transposes under WG semantics)","Upgrade or downgrade JAX to a version whose mosaic_gpu primitives support your transform combination","File an issue with the remaining transforms list if the transform is a legitimate WG-supported case"],"exampleFix":"// before\ninline_mgpu(fn, mgpu_ref.T)\n// after\ninline_mgpu(fn, mgpu_ref)  # transpose handled inside fn manually","handlingStrategy":"validation","validationCode":"assert not user_transforms or all_handled(user_transforms), f'unhandled transforms: {user_transforms}'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Avoid transposes and reshapes on refs passed to inline_mgpu under warp-group semantics","Pin JAX/pallas versions used to develop the kernel"],"tags":["jax","pallas","mosaic-gpu","transforms","not-implemented"],"backgroundTag":"unsupported-transform-combination","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}