{"record":{"id":"7a8f63147f4bdf55","repo":"jax-ml/jax","slug":"cannot-transpose-x-to-self-permutation","errorCode":null,"errorMessage":"Cannot transpose {x} to {self.permutation}","messagePattern":"Cannot transpose (.+?) to (.+?)","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/state/types.py","lineNumber":228,"sourceCode":"    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)\nclass SelectTransform(MultiRefTransform):\n  idx: Array | int\n\n  def transform_types(self, xs):\n    def _type(ref):\n      match ref:\n        case AbstractRef():\n          return ref\n        case core.ShapedArray():\n          raise NotImplementedError","sourceCodeStart":210,"sourceCodeEnd":246,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/state/types.py#L210-L246","documentation":"Raised by TransposeTransform.transform_type when the value being transformed is neither an AbstractRef nor a core.ShapedArray — i.e. the transform encounters an aval type it can't transpose. This is an internal type-dispatch fallback inside JAX's state types.","triggerScenarios":"Applying TransposeTransform to an aval that is a token, abstract unit, or other non-array, non-ref abstract value; usually triggered by composing transforms in unsupported ways.","commonSituations":"Custom primitives or intermediate-language manipulation where non-array avals flow through ref transforms; rarely hit by end users.","solutions":["Inspect the aval type (type(x)) reaching the transform and avoid routing non-array values through it","Ensure only ShapedArray/AbstractRef avals are passed to transpose transforms","Report upstream if it occurs with plain array code"],"exampleFix":null,"handlingStrategy":"type-guard","validationCode":null,"typeGuard":"def is_transposable_aval(x):\n    return isinstance(x, (core.ShapedArray,)) or type(x).__name__ == \"AbstractRef\"","tryCatchPattern":"try:\n    t.transform_type(x)\nexcept TypeError as e:\n    if \"Cannot transpose\" in str(e):\n        raise ValueError(f\"Unsupported aval {type(x)} for transpose\") from e\n    raise","preventionTips":["Only route array/ref avals through transpose transforms","Log aval types when building custom primitives"],"tags":["jax","internal","type-dispatch"],"backgroundTag":"unsupported-operation-type","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}