Skip to content

Commit 9b25e0e

Browse files
Arm backend: Add docstrings for operator_support/right_shift_support.py (#14430)
Signed-off-by: Sebastian Larsson <[email protected]>
1 parent f801c4f commit 9b25e0e

File tree

1 file changed

+16
-1
lines changed

1 file changed

+16
-1
lines changed

backends/arm/operator_support/right_shift_support.py

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,6 +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+
"""Declare operator support for bitwise right-shift in TOSA.
6+
7+
Provide support checks for ``aten.bitwise_right_shift`` and ``__rshift__``
8+
targets across integer and float TOSA profiles.
9+
10+
"""
511

612
# pyre-unsafe
713

@@ -21,6 +27,8 @@
2127

2228
@register_tosa_support_check
2329
class RightShiftSupported(SupportedTOSAOperatorCheck):
30+
"""Provide TOSA support check for right-shift operations."""
31+
2432
targets = [
2533
exir_ops.edge.aten.bitwise_right_shift.Tensor,
2634
exir_ops.edge.aten.__rshift__.Scalar,
@@ -31,8 +39,15 @@ class RightShiftSupported(SupportedTOSAOperatorCheck):
3139
TosaSpecification.create_from_string("TOSA-1.0+FP"),
3240
]
3341

34-
def is_node_tosa_supported(self, node: fx.Node, tosa_spec: TosaSpecification):
42+
def is_node_tosa_supported(
43+
self, node: fx.Node, tosa_spec: TosaSpecification
44+
) -> bool:
45+
"""Return True if the node is supported by TOSA.
46+
47+
Emit a warning on U55 subsets where one-off errors may occur. Otherwise
48+
accept all matching targets.
3549
50+
"""
3651
# TODO MLETORCH-525 Remove warning
3752
if tosa_spec.is_U55_subset:
3853
logging.warning(f"{node.target} may introduce one-off errors.")

0 commit comments

Comments
 (0)