Function _annotate_dropout

backends/example/example_operators/dropout.py:17–34

This is what the graph of a simple clone op looks like: fn_weight = self.fn_weight fn_bias = self.fn_bias permute_copy = torch.ops.aten.permute_copy.default(fn_weight, [1, 0]); fn_weight = None addmm = torch.ops.aten.addmm.default(fn_bias, arg2_1, permute_copy); fn_bias = arg2

(partitions, quant_config)

def _annotate_dropout(partitions, quant_config):
    """
    This is what the graph of a simple clone op looks like:
    fn_weight = self.fn_weight
    fn_bias = self.fn_bias
    permute_copy = torch.ops.aten.permute_copy.default(fn_weight, [1, 0]);  fn_weight = None
    addmm = torch.ops.aten.addmm.default(fn_bias, arg2_1, permute_copy);  fn_bias = arg2_1 = permute_copy = None
    """
    dropout_node = partitions[0].output_nodes[0]
    input_node = dropout_node.args[0]

    if _nodes_are_annotated([dropout_node]):
        return

    _annotate_nodes(
        [(dropout_node, input_node)], quant_config.input_quant_spec, input_node=True
    )
    _annotate_nodes([(dropout_node,)], quant_config.output_quant_spec)


@dataclass