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 dataclasses import dataclass 8 9import torch 10from executorch.backends.example.example_operators.op_base import OpBase 11from executorch.backends.example.example_operators.utils import ( 12 _annotate_nodes, 13 _nodes_are_annotated, 14) 15 16 17def _annotate_dropout(partitions, quant_config): 18 """ 19 This is what the graph of a simple clone op looks like: 20 fn_weight = self.fn_weight 21 fn_bias = self.fn_bias 22 permute_copy = torch.ops.aten.permute_copy.default(fn_weight, [1, 0]); fn_weight = None 23 addmm = torch.ops.aten.addmm.default(fn_bias, arg2_1, permute_copy); fn_bias = arg2_1 = permute_copy = None 24 """ 25 dropout_node = partitions[0].output_nodes[0] 26 input_node = dropout_node.args[0] 27 28 if _nodes_are_annotated([dropout_node]): 29 return 30 31 _annotate_nodes( 32 [(dropout_node, input_node)], quant_config.input_quant_spec, input_node=True 33 ) 34 _annotate_nodes([(dropout_node,)], quant_config.output_quant_spec) 35 36 37@dataclass 38class DropOutNode(OpBase): 39 def __init__(self): 40 super().__init__( 41 # pattern=(torch.clone,), 42 pattern=(torch.nn.modules.dropout.Dropout,), 43 annotate_handle=_annotate_dropout, 44 ) 45