jax-ml/jax · error · ValueError

Permutation {self.permutation} does not match the rank of th

Error message

Permutation {self.permutation} does not match the rank of the type ({x.ndim})

What it means

Raised by TransposeTransform.transform_type when the permutation's length doesn't equal the number of dimensions (ndim) of the ShapedArray being transformed. JAX's state/dispatch transforms use permutations to reorder ref axes, and the permutation must reference every axis exactly once.

Source

Thrown at jax/_src/state/types.py:218

    inverse[p] = i
  return tuple(inverse)


@tree_util.register_dataclass
@dataclasses.dataclass(frozen=True, slots=True)
class TransposeTransform(Transform):
  permutation: tuple[int, ...] = tree.static()

  def undo(self, x: core.AbstractValue) -> Transform:
    return TransposeTransform(_perm_inverse(self.permutation))

  def transform_type(self, x):
    match x:
      case AbstractRef():
        return x.update(inner_aval=self.transform_type(x.inner_aval))
      case core.ShapedArray():
        if len(self.permutation) != x.ndim:
          raise ValueError(
              f"Permutation {self.permutation} does not match the rank of the "
              f"type ({x.ndim})"
          )
        # If there are no explicit axes, do nothing.
        if not all(p is None for p in x.sharding.spec):
          raise NotImplementedError
        new_shape = tuple(x.shape[i] for i in self.permutation)
        return x.update(shape=new_shape)
      case _:
        raise TypeError(f"Cannot transpose {x} to {self.permutation}")

  def pretty_print(self, context: core.JaxprPpContext) -> pp.Doc:
    del context  # Unused.
    return pp.text(f"{{transpose({list(self.permutation)})}}")


@tree_util.register_dataclass
@dataclasses.dataclass(frozen=True, slots=True)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Check len(permutation) == ref.ndim before calling transpose
2. Build the permutation dynamically: perm = tuple(range(ndim)[::-1]) or via numpy argsort
3. Re-derive the permutation after any reshape/slice that changes rank

Example fix

// before
ref.transpose((1, 0))  # ref is 3D
// after
perm = tuple(range(ref.ndim - 1, -1, -1))
ref.transpose(perm)

assert len(perm) == ref.ndim, f"perm {perm} vs ndim {ref.ndim}"

When it happens

Trigger: Calling ref.transpose(...) (or .T on a multi-dim ref) with a permutation tuple whose length differs from ref.ndim; e.g. a 3D ref transposed with (1,0).

Common situations: Passing a hard-coded permutation after changing array rank; slicing a ref (which can drop dims) before transposing; copy-pasting transpose code between tensors of different rank.

Related errors

• unexpected JAX type (e.g. shape/dtype) for argument to VJP f
• cotangent type does not match function output, expected {out
• Pure callbacks do not support transpose. Please use `jax.cus
• Mismatched number of outputs from callback. Expected: {}, Ac
• Incorrect output shape for return value #{i}: Expected: {out

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