{"record":{"id":"64b8b43ef0358800","repo":"jax-ml/jax","slug":"permutation-self-permutation-does-not-match-the","errorCode":null,"errorMessage":"Permutation {self.permutation} does not match the rank of the type ({x.ndim})","messagePattern":"Permutation (.+?) does not match the rank of the type \\((.+?)\\)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/state/types.py","lineNumber":218,"sourceCode":"    inverse[p] = i\n  return tuple(inverse)\n\n\n@tree_util.register_dataclass\n@dataclasses.dataclass(frozen=True, slots=True)\nclass TransposeTransform(Transform):\n  permutation: tuple[int, ...] = tree.static()\n\n  def undo(self, x: core.AbstractValue) -> Transform:\n    return TransposeTransform(_perm_inverse(self.permutation))\n\n  def transform_type(self, x):\n    match x:\n      case AbstractRef():\n        return x.update(inner_aval=self.transform_type(x.inner_aval))\n      case core.ShapedArray():\n        if len(self.permutation) != x.ndim:\n          raise ValueError(\n              f\"Permutation {self.permutation} does not match the rank of the \"\n              f\"type ({x.ndim})\"\n          )\n        # If there are no explicit axes, do nothing.\n        if not all(p is None for p in x.sharding.spec):\n          raise NotImplementedError\n        new_shape = tuple(x.shape[i] for i in self.permutation)\n        return x.update(shape=new_shape)\n      case _:\n        raise TypeError(f\"Cannot transpose {x} to {self.permutation}\")\n\n  def pretty_print(self, context: core.JaxprPpContext) -> pp.Doc:\n    del context  # Unused.\n    return pp.text(f\"{{transpose({list(self.permutation)})}}\")\n\n\n@tree_util.register_dataclass\n@dataclasses.dataclass(frozen=True, slots=True)","sourceCodeStart":200,"sourceCodeEnd":236,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/state/types.py#L200-L236","documentation":"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.","triggerScenarios":"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).","commonSituations":"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.","solutions":["Check len(permutation) == ref.ndim before calling transpose","Build the permutation dynamically: perm = tuple(range(ndim)[::-1]) or via numpy argsort","Re-derive the permutation after any reshape/slice that changes rank"],"exampleFix":"// before\nref.transpose((1, 0))  # ref is 3D\n// after\nperm = tuple(range(ref.ndim - 1, -1, -1))\nref.transpose(perm)","handlingStrategy":"validation","validationCode":"assert len(perm) == ref.ndim, f\"perm {perm} vs ndim {ref.ndim}\"","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Derive permutations from ref.ndim instead of hard-coding","Re-check permutations after reshapes or slices"],"tags":["jax","transpose","shape-mismatch"],"backgroundTag":"shape-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}