xref: /aosp_15_r20/external/executorch/exir/passes/normalize_transpose_pass.py (revision 523fa7a60841cd1ecfb9cc4201f1ca8b03ed023a) 
1# Copyright (c) Meta Platforms, Inc. and affiliates.
2# All rights reserved.
3#
4# This source code is licensed under the BSD-style license found in the
5# LICENSE file in the root directory of this source tree.
6
7import torch
8from executorch.exir.pass_base import ExportPass
9
10
11class NormalizeTransposePass(ExportPass):
12    """
13    Even with functionalization on, we still get graph with
14    torch.ops.aten.t.default op. Ideally we should fix functionalization.
15    TODO: once we have that, we should remove this pass.
16    Check test_normalize_transpose_op in test_passes.py for more details
17    """
18
19    def call_operator(self, op, args, kwargs, meta):
20        if op == torch.ops.aten.t.default:
21            return super().call_operator(
22                torch.ops.aten.t_copy.default, (args[0],), kwargs, meta
23            )
24        return super().call_operator(op, args, kwargs, meta)
25