jax-ml/jax · error · NotImplementedError

Cannot permute last two dimensions with leading dimensions.

Error message

Cannot permute last two dimensions with leading dimensions.

What it means

Raised by the fuser's transpose usage rule: only permutations that keep the last two dimensions in the last two positions (possibly swapped) are supported. The check set(permutation[-2:]) != {permutation[-1], permutation[-2]} catches permutations where the trailing dims were moved into leading positions or vice versa, which the block-index machinery cannot express.

Source

Thrown at jax/_src/pallas/fuser/block_spec.py:2074

  new_permutation = [p for p in permuted_block_dims if p is not None]
  return jax.lax.transpose(x, permutation=new_permutation)


@register_pull_block_spec_rule(lax.transpose_p)
def _transpose_pull_rule(
    ctx: PullRuleContext,
    block_transform: BlockIndexTransform,
    *,
    permutation: tuple[int, ...],
):

  block_shape = block_transform.block_shape
  new_shape = tuple(block_shape[i] for i in permutation)
  aval_in = ctx.avals_in[0]
  assert isinstance(aval_in, core.ShapedArray)
  assert len(block_shape) == len(aval_in.shape)
  if set(permutation[-2:]) != {permutation[-1], permutation[-2]}:
    raise NotImplementedError(
        'Cannot permute last two dimensions with leading dimensions.'
    )

  def new_block_index_transform(*idxs):
    original_idxs = block_transform.block_index_transform(*idxs)
    return tuple(original_idxs[i] for i in permutation)

  return [block_transform.replace(
      block_shape=new_shape,
      block_index_transform=new_block_index_transform)]


@register_eval_rule(lax.split_p)
def _split_eval_rule(
    eval_ctx: KernelEvalContext, x, sizes: Sequence[int], axis: int
):
  del eval_ctx
  return lax.split(x, sizes=sizes, axis=axis)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Restructure to only swap the last two axes (single lax.transpose((..., 1, 0)) on the tail) and handle leading-axis reordering via block index maps or separate reshapes
  2. Move the general transpose outside the fused kernel and pass the pre-transposed array in
  3. Express leading-dim permutations through grid/index_map instead of an in-kernel transpose

Example fix

// before
y = jnp.transpose(x, (2, 0, 1))  # moves last dim to front: unsupported
// after
y = jnp.swapaxes(x, -1, -2)  # only permutes last two dims: supported
Defensive patterns

Strategy: validation

Validate before calling

def transpose_supported(perm) -> bool:
    perm = list(perm)
    return set(perm[-2:]) == {perm[-1], perm[-2]}
assert transpose_supported(perm), 'only last-two-dims permutation supported'

Type guard

def is_tail_only_perm(perm: tuple[int, ...]) -> bool:
    return sorted(perm[-2:]) == [len(perm) - 2, len(perm) - 1]

Try / catch

try:
    y = fused_t(x, perm)
except NotImplementedError as e:
    if 'Cannot permute last two' in str(e):
        y = jnp.transpose(x, perm)  # outside fusion
    else:
        raise

Prevention

When it happens

Trigger: Calling lax.transpose / jnp.transpose inside a fused Pallas region with a permutation that moves one of the last two axes to a leading position, e.g. jnp.transpose(x, (2, 0, 1)) for a 3-D operand.

Common situations: General N-D transposes (like NCHW<->NHWC with more than 2 trailing dims, or batched matrix transposes mixing batch axes with matrix axes); refactoring existing kernels that reshape+transpose for matmul layout; JAX version upgrades tightening transpose support in the fuser.

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

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