jax-ml/jax · error · NotImplementedError

Encountered non-leading UntilingTransform or UnswizzleRef tr

Error message

Encountered non-leading UntilingTransform or UnswizzleRef transforms: {transforms}

What it means

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.

Source

Thrown at jax/_src/pallas/mosaic_gpu/lowering.py:4773

      ref, transforms = ref
      ref_aval, transform_avals = ref_aval
      # We ignore other transforms here, because they are already embedded
      # in the jaxpr.
      assert isinstance(ref_aval, state_types.AbstractRef)
      ref, ref_aval, _ = _handle_transforms(
          ctx, ref_aval, ref, transform_avals, transforms,
          handle_reshapes=False, handle_transposes=False
      )
      if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Warpgroup:
        # In warpgroup semantics, we must reapply the transforms that were on
        # the `BlockSpec` here, as the below expects the transformed value to be
        # fed in.
        spec_transforms = tuple(
            t for t in transforms
            if isinstance(t, (gpu_core.UntilingTransform, gpu_core.UnswizzleRef))
        )
        if spec_transforms != transforms[:len(spec_transforms)]:
          raise NotImplementedError(
              "Encountered non-leading UntilingTransform or UnswizzleRef "
              f"transforms: {transforms}"
          )
        for t in pallas_core.undo_transforms(ref_aval, spec_transforms):
          ref_aval = cast(state_types.AbstractRef, t.transform_type(ref_aval))
        ref = _reinterpret_cast(ref, ref_aval)
    args.append(ref)
  program_ids = program_ids_treedef.unflatten(flat_program_ids)
  for axis, pid in enumerate(program_ids):
    if pid is not None:
      continue
    program_ids[axis] = _program_id(
        axis, ctx.module_ctx.squashed_dims, len(program_ids)
    )
  new_module_ctx = dataclasses.replace(ctx.module_ctx, program_ids=program_ids)
  return lower_jaxpr_to_mosaic_gpu(
      new_module_ctx, ctx.launch_ctx, jaxpr, args
  )

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Upgrade JAX — transform ordering bugs are fixed upstream over time
2. Simplify the kernel: avoid combining swizzled layouts with discharge/untile on the same ref
3. Report a minimal repro to the JAX team

When it happens

Trigger: 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.

Common situations: Using swizzled buffers combined with multiple transforms (scatter/discharge/untile) in complex pallas pipelines; often a JAX-internal ordering issue rather than user-controllable.

Understand the failure class

Background: UnsupportedOperationException and "is not supported" errors: when a library deliberately refuses a call — this error's family across 30 libraries.

Related errors

• Unsupported transform: {type(transform)}
• Non-indexing transforms on GMEM refs are not implemented.
• Not all transforms could be handled. Remaining transforms: {
• Unsupported dtype: {ref.dtype}
• Only SMEM and TMEM refs are supported.

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/48c0c7b29e1156cc. Report an issue: GitHub.