xref: /aosp_15_r20/external/executorch/backends/example/example_operators/ops.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
7from collections import OrderedDict
8
9from executorch.backends.example.example_operators.adaptive_avg_pool2d import (
10    AdaptiveAvgPool2dNode,
11)
12
13from executorch.backends.example.example_operators.add import AddNode
14from executorch.backends.example.example_operators.conv2d import Conv2DNode
15from executorch.backends.example.example_operators.conv_relu import ConvReluNode
16from executorch.backends.example.example_operators.dropout import DropOutNode
17from executorch.backends.example.example_operators.flatten import FlattenNode
18from executorch.backends.example.example_operators.linear import LinearNode
19
20# The ordering of this is important as the quantizer will try to match the patterns in this order.
21# That's why we want to match the fused patterns first and then the non-fused ones.
22module_to_annotator = OrderedDict(
23    {
24        ConvReluNode().pattern: ConvReluNode(),
25        Conv2DNode().pattern: Conv2DNode(),
26        LinearNode().pattern: LinearNode(),
27        AddNode().pattern: AddNode(),
28        AdaptiveAvgPool2dNode().pattern: AdaptiveAvgPool2dNode(),
29        FlattenNode().pattern: FlattenNode(),
30        DropOutNode().pattern: DropOutNode(),
31    }
32)
33