{"record":{"id":"48c0c7b29e1156cc","repo":"jax-ml/jax","slug":"encountered-non-leading-untilingtransform-or-unswi","errorCode":null,"errorMessage":"Encountered non-leading UntilingTransform or UnswizzleRef transforms: {transforms}","messagePattern":"Encountered non-leading UntilingTransform or UnswizzleRef transforms: (.+?)","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/lowering.py","lineNumber":4773,"sourceCode":"      ref, transforms = ref\n      ref_aval, transform_avals = ref_aval\n      # We ignore other transforms here, because they are already embedded\n      # in the jaxpr.\n      assert isinstance(ref_aval, state_types.AbstractRef)\n      ref, ref_aval, _ = _handle_transforms(\n          ctx, ref_aval, ref, transform_avals, transforms,\n          handle_reshapes=False, handle_transposes=False\n      )\n      if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Warpgroup:\n        # In warpgroup semantics, we must reapply the transforms that were on\n        # the `BlockSpec` here, as the below expects the transformed value to be\n        # fed in.\n        spec_transforms = tuple(\n            t for t in transforms\n            if isinstance(t, (gpu_core.UntilingTransform, gpu_core.UnswizzleRef))\n        )\n        if spec_transforms != transforms[:len(spec_transforms)]:\n          raise NotImplementedError(\n              \"Encountered non-leading UntilingTransform or UnswizzleRef \"\n              f\"transforms: {transforms}\"\n          )\n        for t in pallas_core.undo_transforms(ref_aval, spec_transforms):\n          ref_aval = cast(state_types.AbstractRef, t.transform_type(ref_aval))\n        ref = _reinterpret_cast(ref, ref_aval)\n    args.append(ref)\n  program_ids = program_ids_treedef.unflatten(flat_program_ids)\n  for axis, pid in enumerate(program_ids):\n    if pid is not None:\n      continue\n    program_ids[axis] = _program_id(\n        axis, ctx.module_ctx.squashed_dims, len(program_ids)\n    )\n  new_module_ctx = dataclasses.replace(ctx.module_ctx, program_ids=program_ids)\n  return lower_jaxpr_to_mosaic_gpu(\n      new_module_ctx, ctx.launch_ctx, jaxpr, args\n  )","sourceCodeStart":4755,"sourceCodeEnd":4791,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/lowering.py#L4755-L4791","documentation":"During transform handling for references, UntilingTransform/UnswizzleRef transforms must appear at the front of the transform list. If they occur after other transforms, this NotImplementedError fires because undoing them out of order is not supported.","triggerScenarios":"A ref carrying a mixed transform stack where un-tiling/un-swizzling is not leading, e.g. discharge followed by unswizzle ordering produced by nested transforms on buffers.","commonSituations":"Using swizzled buffers combined with multiple transforms (scatter/discharge/untile) in complex pallas pipelines; often a JAX-internal ordering issue rather than user-controllable.","solutions":["Upgrade JAX — transform ordering bugs are fixed upstream over time","Simplify the kernel: avoid combining swizzled layouts with discharge/untile on the same ref","Report a minimal repro to the JAX team"],"exampleFix":null,"handlingStrategy":"retry","validationCode":null,"typeGuard":null,"tryCatchPattern":null,"preventionTips":["Keep JAX current — transform-ordering fixes land upstream","Avoid mixing swizzle and discharge transforms on one ref"],"tags":["pallas","mosaic-gpu","transforms","not-implemented"],"backgroundTag":"unsupported-operation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}