6
6
# LICENSE file in the root directory of this source tree.
7
7
8
8
# pyre-unsafe
9
+ """Provide utilities for quantization annotations.
9
10
10
- #
11
- # Utility functions for TOSAQuantizer
12
- #
11
+ Use these helpers to check and mark annotation state when working with
12
+ ``QuantizationAnnotation`` entries in FX node metadata.
13
+
14
+ """
13
15
14
16
from typing import cast
15
17
20
22
21
23
22
24
def is_annotated (node : Node ) -> bool :
23
- """Given a node return whether the node is annotated."""
25
+ """Return True if the node is annotated.
26
+
27
+ Args:
28
+ node (Node): FX node to inspect.
29
+
30
+ Returns:
31
+ bool: True if ``Q_ANNOTATION_KEY`` exists and ``_annotated`` is set.
32
+
33
+ """
24
34
return (
25
35
Q_ANNOTATION_KEY in node .meta
26
36
and cast (QuantizationAnnotation , node .meta [Q_ANNOTATION_KEY ])._annotated
27
37
)
28
38
29
39
30
40
def is_output_annotated (node : Node ) -> bool :
31
- """Given a node, return whether the output of the node is annotated."""
41
+ """Return True if the node's output is annotated.
42
+
43
+ Args:
44
+ node (Node): FX node to inspect.
45
+
46
+ Returns:
47
+ bool: True if annotated and an output qspec is present.
48
+
49
+ """
32
50
if Q_ANNOTATION_KEY in node .meta :
33
51
annotation = cast (QuantizationAnnotation , node .meta [Q_ANNOTATION_KEY ])
34
52
return annotation ._annotated and annotation .output_qspec is not None
@@ -37,8 +55,14 @@ def is_output_annotated(node: Node) -> bool:
37
55
38
56
39
57
def mark_node_as_annotated (node : Node ) -> None :
40
- """Marks node as annotated. If needed, an empty QuantizationAnnotation is added
41
- to the quantization_annotation node meta entry.
58
+ """Mark a node as annotated.
59
+
60
+ Create an empty ``QuantizationAnnotation`` on the node when missing and set
61
+ its ``_annotated`` flag to True.
62
+
63
+ Args:
64
+ node (Node): FX node to update.
65
+
42
66
"""
43
67
if Q_ANNOTATION_KEY not in node .meta :
44
68
node .meta [Q_ANNOTATION_KEY ] = QuantizationAnnotation ()
0 commit comments