Function _annotate_linear

backends/example/example_operators/linear.py:17–53

This is what the graph of a simple linear 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 = arg

(partitions, quant_config)

def _annotate_linear(partitions, quant_config):
    """
    This is what the graph of a simple linear 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
    """
    linear_node = partitions[0].output_nodes[0]
    if _nodes_are_annotated([linear_node]):
        return

    input_node = linear_node.args[0]
    # permute_node = linear_node.args[1]
    # print("permute_node: ", permute_node, " args: ", permute_node.args, " target: ", permute_node.target)
    weight_node = linear_node.args[1]
    print(
        "weight_node: ",
        weight_node,
        " args: ",
        weight_node.args,
        " target: ",
        weight_node.target,
    )
    # Unused.
    # bias_node = output_node.args[0]

    # if _nodes_are_annotated([linear_node, permute_node]):
    #     return

    _annotate_nodes(
        [(linear_node, input_node)], quant_config.input_quant_spec, input_node=True
    )
    _annotate_nodes(
        [(linear_node, weight_node)], quant_config.weight_quant_spec, input_node=True
    )
    _annotate_nodes([(linear_node,)], quant_config.output_quant_spec)


@dataclass