{"record":{"id":"651e23324bd3db8b","repo":"jax-ml/jax","slug":"transform-mismatch-got-user-transforms-expecte","errorCode":null,"errorMessage":"Transform mismatch: got {user_transforms}, expected {ty_transforms}","messagePattern":"Transform mismatch: got (.+?), expected (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/primitives.py","lineNumber":3624,"sourceCode":"        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\n\n@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,","sourceCodeStart":3606,"sourceCodeEnd":3642,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/primitives.py#L3606-L3642","documentation":"Raised when the transforms explicitly supplied by the user for an inline_mgpu argument do not match the transforms derived from the argument's abstract value (aval) via pallas_core.undo_transforms. This is a consistency check: the declared transform chain and the actual aval's baked-in transforms diverge.","triggerScenarios":"Calling inline_mgpu with user_transforms that differ in content or order from tuple(pallas_core.undo_transforms(aval, t.transforms)); typically from manually constructing transformed refs or stale cached avals.","commonSituations":"Mixing transform application order; constructing refs through internal APIs rather than public helpers; JAX version change altering transform normalization order.","solutions":["Pass the exact transforms reported in 'expected' (ty_transforms) instead of hand-built ones","Rebuild the ref from scratch with public APIs so transforms are attached automatically","Clear JIT caches and retry to rule out stale avals","Align JAX/pallas versions between components generating and consuming the transforms"],"exampleFix":"// before\ninline_mgpu(fn, ref, transforms=(Transpose(0,1),))\n// after\ninline_mgpu(fn, ref, transforms=tuple(pallas_core.undo_transforms(aval, ref.transforms)))","handlingStrategy":"validation","validationCode":"expected = tuple(pallas_core.undo_transforms(aval, t.transforms))\nassert tuple(user_transforms) == expected, (user_transforms, expected)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Never hand-construct transform tuples; derive them from the ref","Clear caches after changing kernel signatures"],"tags":["jax","pallas","transform-mismatch","validation"],"backgroundTag":"argument-mismatch-validation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}