{"record":{"id":"ddb6f8adbac02877","repo":"jax-ml/jax","slug":"cannot-permute-last-two-dimensions-with-leading-di","errorCode":null,"errorMessage":"Cannot permute last two dimensions with leading dimensions.","messagePattern":"Cannot permute last two dimensions with leading dimensions\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/fuser/block_spec.py","lineNumber":2074,"sourceCode":"  new_permutation = [p for p in permuted_block_dims if p is not None]\n  return jax.lax.transpose(x, permutation=new_permutation)\n\n\n@register_pull_block_spec_rule(lax.transpose_p)\ndef _transpose_pull_rule(\n    ctx: PullRuleContext,\n    block_transform: BlockIndexTransform,\n    *,\n    permutation: tuple[int, ...],\n):\n\n  block_shape = block_transform.block_shape\n  new_shape = tuple(block_shape[i] for i in permutation)\n  aval_in = ctx.avals_in[0]\n  assert isinstance(aval_in, core.ShapedArray)\n  assert len(block_shape) == len(aval_in.shape)\n  if set(permutation[-2:]) != {permutation[-1], permutation[-2]}:\n    raise NotImplementedError(\n        'Cannot permute last two dimensions with leading dimensions.'\n    )\n\n  def new_block_index_transform(*idxs):\n    original_idxs = block_transform.block_index_transform(*idxs)\n    return tuple(original_idxs[i] for i in permutation)\n\n  return [block_transform.replace(\n      block_shape=new_shape,\n      block_index_transform=new_block_index_transform)]\n\n\n@register_eval_rule(lax.split_p)\ndef _split_eval_rule(\n    eval_ctx: KernelEvalContext, x, sizes: Sequence[int], axis: int\n):\n  del eval_ctx\n  return lax.split(x, sizes=sizes, axis=axis)","sourceCodeStart":2056,"sourceCodeEnd":2092,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/fuser/block_spec.py#L2056-L2092","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["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","Move the general transpose outside the fused kernel and pass the pre-transposed array in","Express leading-dim permutations through grid/index_map instead of an in-kernel transpose"],"exampleFix":"// before\ny = jnp.transpose(x, (2, 0, 1))  # moves last dim to front: unsupported\n// after\ny = jnp.swapaxes(x, -1, -2)  # only permutes last two dims: supported\n","handlingStrategy":"validation","validationCode":"def transpose_supported(perm) -> bool:\n    perm = list(perm)\n    return set(perm[-2:]) == {perm[-1], perm[-2]}\nassert transpose_supported(perm), 'only last-two-dims permutation supported'","typeGuard":"def is_tail_only_perm(perm: tuple[int, ...]) -> bool:\n    return sorted(perm[-2:]) == [len(perm) - 2, len(perm) - 1]","tryCatchPattern":"try:\n    y = fused_t(x, perm)\nexcept NotImplementedError as e:\n    if 'Cannot permute last two' in str(e):\n        y = jnp.transpose(x, perm)  # outside fusion\n    else:\n        raise","preventionTips":["Restrict in-kernel transposes to swapping the last two axes","Move N-D transposes out of the fused region","Encode leading-dim reordering in index_map instead"],"tags":["jax","pallas","transpose","permutation","not-implemented"],"backgroundTag":"unsupported-operation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}