jax-ml/jax · error · TypeError

Cannot transpose {x} to {self.permutation}

Error message

Cannot transpose {x} to {self.permutation}

What it means

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.

Source

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

    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)
class SelectTransform(MultiRefTransform):
  idx: Array | int

  def transform_types(self, xs):
    def _type(ref):
      match ref:
        case AbstractRef():
          return ref
        case core.ShapedArray():
          raise NotImplementedError

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Inspect the aval type (type(x)) reaching the transform and avoid routing non-array values through it
2. Ensure only ShapedArray/AbstractRef avals are passed to transpose transforms
3. Report upstream if it occurs with plain array code

def is_transposable_aval(x):
    return isinstance(x, (core.ShapedArray,)) or type(x).__name__ == "AbstractRef"

try:
    t.transform_type(x)
except TypeError as e:
    if "Cannot transpose" in str(e):
        raise ValueError(f"Unsupported aval {type(x)} for transpose") from e
    raise

When it happens

Trigger: 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.

Common situations: Custom primitives or intermediate-language manipulation where non-array avals flow through ref transforms; rarely hit by end users.

Related errors

• scan number of arguments doesn't match the number of jaxpr a
• Unrecognized index type: {typ}
• static_slice: unrecognized index {pidx.index} at position {p
• static_slice: unrecognized index {pidx.index}
• dynamic_slice: unrecognized index {pidx.index}

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