Skip to content

Commit 6ec3149

Browse files
Arm backend: Add docstrings for operator_support/reduce_sum_support.py (#15638)
Signed-off-by: Sebastian Larsson <[email protected]>
1 parent 80c9040 commit 6ec3149

File tree

1 file changed

+16
-1
lines changed

1 file changed

+16
-1
lines changed

backends/arm/operator_support/reduce_sum_support.py

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,11 @@
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+
"""Declare operator support for ``aten.sum.dim_IntList`` in TOSA.
56
7+
Provide shape constraints for U55 subsets; otherwise allow reductions.
8+
9+
"""
610
from typing import cast
711

812
import torch.fx as fx
@@ -16,14 +20,25 @@
1620

1721
@register_tosa_support_check
1822
class SumSupported(SupportedTOSAOperatorCheck):
23+
"""Provide TOSA support check for sum over dimensions."""
24+
1925
targets = [exir_ops.edge.aten.sum.dim_IntList]
2026

2127
tosa_specs = [
2228
TosaSpecification.create_from_string("TOSA-1.0+INT"),
2329
TosaSpecification.create_from_string("TOSA-1.0+FP"),
2430
]
2531

26-
def is_node_tosa_supported(self, node: fx.Node, tosa_spec: TosaSpecification):
32+
def is_node_tosa_supported(
33+
self, node: fx.Node, tosa_spec: TosaSpecification
34+
) -> bool:
35+
"""Return True if the node is supported by TOSA.
36+
37+
On U55 subsets, enforce bounds on the reduced dimension and the products
38+
of sizes before/after the reduction axis. On other targets, accept the
39+
operation unconditionally.
40+
41+
"""
2742
if not tosa_spec.is_U55_subset:
2843
return True
2944

0 commit comments

Comments
 (0)