Function _annotate_flatten

backends/example/example_operators/flatten.py:17–31

This is what the graph of a simple add op looks like: add_tensor = torch.ops.aten.add.Tensor(arg0_1, arg1_1); arg0_1 = arg1_1 = None

(partitions, quant_config)

def _annotate_flatten(partitions, quant_config):
    """
    This is what the graph of a simple add op looks like:
    add_tensor = torch.ops.aten.add.Tensor(arg0_1, arg1_1);  arg0_1 = arg1_1 = None
    """
    flatten_node = partitions[0].output_nodes[0]
    flatten_input = flatten_node.args[0]

    if _nodes_are_annotated([flatten_node]):
        return

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


@dataclass