From d98148d6b23540d91be33ef8cd0ce27f8ad28e77 Mon Sep 17 00:00:00 2001 From: Sebastian Larsson Date: Thu, 11 Sep 2025 10:12:57 +0200 Subject: [PATCH] Arm backend: Add docstrings for operator_support/reduce_sum_support.py Change-Id: I013f34240d564b371968b16f9e78e367bdf626fe Signed-off-by: Sebastian Larsson --- .../arm/operator_support/reduce_sum_support.py | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/backends/arm/operator_support/reduce_sum_support.py b/backends/arm/operator_support/reduce_sum_support.py index 76d1ba7bf36..02e9e0db90e 100644 --- a/backends/arm/operator_support/reduce_sum_support.py +++ b/backends/arm/operator_support/reduce_sum_support.py @@ -2,7 +2,11 @@ # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +"""Declare operator support for ``aten.sum.dim_IntList`` in TOSA. +Provide shape constraints for U55 subsets; otherwise allow reductions. + +""" from typing import cast import torch.fx as fx @@ -16,6 +20,8 @@ @register_tosa_support_check class SumSupported(SupportedTOSAOperatorCheck): + """Provide TOSA support check for sum over dimensions.""" + targets = [exir_ops.edge.aten.sum.dim_IntList] tosa_specs = [ @@ -23,7 +29,16 @@ class SumSupported(SupportedTOSAOperatorCheck): TosaSpecification.create_from_string("TOSA-1.0+FP"), ] - def is_node_tosa_supported(self, node: fx.Node, tosa_spec: TosaSpecification): + def is_node_tosa_supported( + self, node: fx.Node, tosa_spec: TosaSpecification + ) -> bool: + """Return True if the node is supported by TOSA. + + On U55 subsets, enforce bounds on the reduced dimension and the products + of sizes before/after the reduction axis. On other targets, accept the + operation unconditionally. + + """ if not tosa_spec.is_U55_subset: return True