utils.py

# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

from torch.ao.quantization.quantizer.quantizer import QuantizationAnnotation


def _nodes_are_annotated(node_list):
    for node in node_list:
        quantization_annotation = node.meta.get("quantization_annotation", None)
        if not quantization_annotation:
            return False
        if quantization_annotation._annotated:
            continue
        else:
            return False
    return True


def _annotate_nodes(node_tuples, quant_spec, input_node=False):
    for node_tuple in node_tuples:
        node = node_tuple[0]
        quant_annotation = node.meta.get(
            "quantization_annotation", QuantizationAnnotation(_annotated=True)
        )
        if input_node:
            input_node = node_tuple[1]
            quant_annotation.input_qspec_map[input_node] = quant_spec
        else:
            quant_annotation.output_qspec = quant_spec
        node.meta["quantization_annotation"] = quant_annotation