    def _infer_shapes(self, x: Tensor) -> List[IntVar]:
        """Infers shapes for permute021."""
        x_shape = x._attrs["shape"]
        return x_shape[:-2] + [x_shape[-1], x_shape[-2]]
Infers shapes for permute021.
python/aitemplate/compiler/ops/tensor/permute021.py