From 4751b46470a7bc8e97b9832657354fedaecbc8d0 Mon Sep 17 00:00:00 2001 From: Sebastian Larsson Date: Wed, 10 Sep 2025 15:38:04 +0200 Subject: [PATCH] Arm backend: Add docstrings to init and arm_quantizer_utils in quantizer Signed-off-by: Sebastian Larsson Change-Id: I9d9b451137b62214a9e245879b67567921650047 --- backends/arm/quantizer/__init__.py | 5 +++ backends/arm/quantizer/arm_quantizer_utils.py | 38 +++++++++++++++---- 2 files changed, 36 insertions(+), 7 deletions(-) diff --git a/backends/arm/quantizer/__init__.py b/backends/arm/quantizer/__init__.py index 5cb5c834a98..e36c683416a 100644 --- a/backends/arm/quantizer/__init__.py +++ b/backends/arm/quantizer/__init__.py @@ -2,7 +2,12 @@ # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +"""Expose quantizer APIs and load optional quantized kernels. +Import the public quantizer classes and configuration helpers for Arm +backends. Attempt to load portable and quantized libraries; fall back to a +log message if unavailable. +""" from .quantization_config import QuantizationConfig # noqa # usort: skip from .arm_quantizer import ( # noqa diff --git a/backends/arm/quantizer/arm_quantizer_utils.py b/backends/arm/quantizer/arm_quantizer_utils.py index 838dd44733e..90876386aa6 100644 --- a/backends/arm/quantizer/arm_quantizer_utils.py +++ b/backends/arm/quantizer/arm_quantizer_utils.py @@ -6,10 +6,12 @@ # LICENSE file in the root directory of this source tree. # pyre-unsafe +"""Provide utilities for quantization annotations. -# -# Utility functions for TOSAQuantizer -# +Use these helpers to check and mark annotation state when working with +``QuantizationAnnotation`` entries in FX node metadata. + +""" from typing import cast @@ -20,7 +22,15 @@ def is_annotated(node: Node) -> bool: - """Given a node return whether the node is annotated.""" + """Return True if the node is annotated. + + Args: + node (Node): FX node to inspect. + + Returns: + bool: True if ``Q_ANNOTATION_KEY`` exists and ``_annotated`` is set. + + """ return ( Q_ANNOTATION_KEY in node.meta and cast(QuantizationAnnotation, node.meta[Q_ANNOTATION_KEY])._annotated @@ -28,7 +38,15 @@ def is_annotated(node: Node) -> bool: def is_output_annotated(node: Node) -> bool: - """Given a node, return whether the output of the node is annotated.""" + """Return True if the node's output is annotated. + + Args: + node (Node): FX node to inspect. + + Returns: + bool: True if annotated and an output qspec is present. + + """ if Q_ANNOTATION_KEY in node.meta: annotation = cast(QuantizationAnnotation, node.meta[Q_ANNOTATION_KEY]) return annotation._annotated and annotation.output_qspec is not None @@ -37,8 +55,14 @@ def is_output_annotated(node: Node) -> bool: def mark_node_as_annotated(node: Node) -> None: - """Marks node as annotated. If needed, an empty QuantizationAnnotation is added - to the quantization_annotation node meta entry. + """Mark a node as annotated. + + Create an empty ``QuantizationAnnotation`` on the node when missing and set + its ``_annotated`` flag to True. + + Args: + node (Node): FX node to update. + """ if Q_ANNOTATION_KEY not in node.meta: node.meta[Q_ANNOTATION_KEY] = QuantizationAnnotation()