2
2
#
3
3
# This source code is licensed under the BSD-style license found in the
4
4
# 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
+ """
5
11
6
12
# pyre-unsafe
7
13
21
27
22
28
@register_tosa_support_check
23
29
class RightShiftSupported (SupportedTOSAOperatorCheck ):
30
+ """Provide TOSA support check for right-shift operations."""
31
+
24
32
targets = [
25
33
exir_ops .edge .aten .bitwise_right_shift .Tensor ,
26
34
exir_ops .edge .aten .__rshift__ .Scalar ,
@@ -31,8 +39,15 @@ class RightShiftSupported(SupportedTOSAOperatorCheck):
31
39
TosaSpecification .create_from_string ("TOSA-1.0+FP" ),
32
40
]
33
41
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.
35
49
50
+ """
36
51
# TODO MLETORCH-525 Remove warning
37
52
if tosa_spec .is_U55_subset :
38
53
logging .warning (f"{ node .target } may introduce one-off errors." )
0 commit comments