Skip to content

Commit 9b598be

Browse files
Sebastian-LarssonStrycekSimon
authored andcommitted
Arm backend: Add docstrings to init and arm_quantizer_utils in quantizer (pytorch#14375)
Signed-off-by: Sebastian Larsson <[email protected]>
1 parent 97b9bcb commit 9b598be

File tree

2 files changed

+36
-7
lines changed

2 files changed

+36
-7
lines changed

backends/arm/quantizer/__init__.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,12 @@
22
#
33
# This source code is licensed under the BSD-style license found in the
44
# LICENSE file in the root directory of this source tree.
5+
"""Expose quantizer APIs and load optional quantized kernels.
56
7+
Import the public quantizer classes and configuration helpers for Arm
8+
backends. Attempt to load portable and quantized libraries; fall back to a
9+
log message if unavailable.
10+
"""
611

712
from .quantization_config import QuantizationConfig # noqa # usort: skip
813
from .arm_quantizer import ( # noqa

backends/arm/quantizer/arm_quantizer_utils.py

Lines changed: 31 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -6,10 +6,12 @@
66
# LICENSE file in the root directory of this source tree.
77

88
# pyre-unsafe
9+
"""Provide utilities for quantization annotations.
910
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+
"""
1315

1416
from typing import cast
1517

@@ -20,15 +22,31 @@
2022

2123

2224
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+
"""
2434
return (
2535
Q_ANNOTATION_KEY in node.meta
2636
and cast(QuantizationAnnotation, node.meta[Q_ANNOTATION_KEY])._annotated
2737
)
2838

2939

3040
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+
"""
3250
if Q_ANNOTATION_KEY in node.meta:
3351
annotation = cast(QuantizationAnnotation, node.meta[Q_ANNOTATION_KEY])
3452
return annotation._annotated and annotation.output_qspec is not None
@@ -37,8 +55,14 @@ def is_output_annotated(node: Node) -> bool:
3755

3856

3957
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+
4266
"""
4367
if Q_ANNOTATION_KEY not in node.meta:
4468
node.meta[Q_ANNOTATION_KEY] = QuantizationAnnotation()

0 commit comments

Comments
 (0)