Arm backend: Format Arm passes (group 7) (pytorch#17615)

Apply docformatter to the final set of passes.

Signed-off-by: Sebastian Larsson <sebastian.larsson@arm.com>

Author: Sebastian-Larsson

.lintrunner.toml

@@ -510,7 +510,6 @@ include_patterns = ['backends/arm/**/*.py']
 exclude_patterns = [
   'third-party/**',
   '**/third-party/**',
-  'backends/arm/_passes/**',
   'backends/arm/test/**',
 ]
 command = [

backends/arm/_passes/arm_pass.py

@@ -38,7 +38,9 @@ def allowed_to_transform(self, meta: NodeMetadata | dict[str, Any]) -> bool:
         return not disallow_tfa
 
     def _is_quantized_meta(self, meta: NodeMetadata | dict[str, Any]) -> bool:
-        """Return True when meta indicates fully quantized inputs and outputs."""
+        """Return True when meta indicates fully quantized inputs and
+        outputs.
+        """
         if isinstance(meta, NodeMetadata):
             meta_dict = meta.data
         else:
@@ -107,9 +109,12 @@ def call_submodule(
     def call_shape_operator(
         self, op, args: tuple, kwargs: dict, meta: NodeMetadata, update: bool
     ) -> ProxyValue:
-        """
-        Call operator for shape-producing operators. This function is responsible for marking the output of the operator
-        with the TosaSpecialDtype of SHAPE, so that later passes can identify it as a shape-producing operator and handle it accordingly.
+        """Call operator for shape-producing operators.
+
+        This function is responsible for marking the output of the operator with
+        the TosaSpecialDtype of SHAPE, so that later passes can identify it as a
+        shape-producing operator and handle it accordingly.
+
         """
         # Copy meta and set TosaSpecialDtype to SHAPE
         if not isinstance(meta, NodeMetadata):

backends/arm/_passes/constant_folding_pass.py

@@ -12,7 +12,11 @@
 
 
 class ConstantFoldingPass(ArmPass):
-    """Fold constant subgraphs using torch's export constant folding pass. To be used before to_edge transform."""
+    """Fold constant subgraphs using torch's export constant folding pass.
+
+    To be used before to_edge transform.
+
+    """
 
     _passes_required_after: Set[Type[ExportPass]] = set()
 
@@ -31,9 +35,7 @@ class ConstantFoldingPass(ArmPass):
     def _ensure_param_attr(
         self, graph_module: torch.fx.GraphModule, node: torch.fx.Node | None
     ) -> bool:
-        """
-        Replaces tensor attributes with parameter attributes.
-        """
+        """Replaces tensor attributes with parameter attributes."""
         if node is None or node.op != "get_attr":
             return False
         target = node.target

backends/arm/_passes/control_flow_const_inline.py

@@ -15,11 +15,14 @@
 
 
 class ControlFlowConstInlinePass(ArmPass):
-    """
-        When we lift out each control flow body as its own GraphModule, any scalar constants that were captured in Python become module attributes.  FX represents those as get_attr nodes in the
-    submodule graph. These become getattr nodes submodule graph.
+    """When we lift out each control flow body as its own GraphModule, any
+    scalar constants that were captured in Python become module attributes.  FX
+    represents those as get_attr nodes in the submodule graph. These become
+    getattr nodes submodule graph.
+
+    This pass ensures that Scalar tensors in control flow operation are
+    converted from getattr operators to expected call_function full ops.
 
-        This pass ensures that Scalar tensors in control flow operation are converted from getattr operators to expected call_function full ops.
     """
 
     _passes_required_after: Set[Type[ExportPass]] = set()

backends/arm/_passes/convert_full_like_to_full_pass.py

@@ -20,15 +20,15 @@ class ConvertFullLikeToFullPass(ArmPass):
     As per the full_like PyTorch documentation, `torch.full_like(input,
     fill_value)` is equivalent to:
 
-    ```
-    torch.full(
-        input.size(),
-        fill_value,
-        dtype=input.dtype,
-        layout=input.layout,
-        device=input.device,
-    )
-    ```
+    ::
+
+        torch.full(
+            input.size(),
+            fill_value,
+            dtype=input.dtype,
+            layout=input.layout,
+            device=input.device,
+        )
 
     Skip layout and device since it's not relevant for our backend.
 

backends/arm/_passes/decompose_strided_slice_copy_pass.py

@@ -43,8 +43,8 @@ def _fixup_end(end, dim_size):
 
 
 class DecomposeStridedSliceCopyPass(ArmPass):
-    """
-    Decompose edge.aten.slice_copy.Tensor with non-unit step into supported ops.
+    """Decompose edge.aten.slice_copy.Tensor with non-unit step into supported
+    ops.
 
     Given:
         out = slice_copy(x, dim, start, end, step)   with step > 1
@@ -57,6 +57,7 @@ class DecomposeStridedSliceCopyPass(ArmPass):
       5) out = view_copy(y3, ...)                    # collapse the singleton dim
 
     This implements "take every step-th element" using only unit-step slice + reshape.
+
     """
 
     _passes_required_after: Set[Type[ExportPass]] = set()

backends/arm/_passes/decompose_tril_pass.py

@@ -45,9 +45,12 @@ def _get_ops(op):
 
 
 class DecomposeTrilPass(ArmPass):
-    """
-    mask_bool = (row + diagonal) >= col     (intended AOT-constant)
-    out = where(mask_bool, x, 0)            (0 is a scalar tensor, broadcasted)
+    """Tril decomposition.
+
+    Decomposition:
+        mask_bool = (row + diagonal) >= col  (intended AOT-constant)
+        out = where(mask_bool, x, 0)         (0 is a scalar tensor, broadcasted)
+
     """
 
     _passes_required_after: Set[Type[ExportPass]] = {ComputeConstantOpsAOTPass}

backends/arm/_passes/scalars_to_attribute_pass.py

@@ -41,8 +41,7 @@ def _convert_scalar_args(
         graph_module: GraphModule,
         n: Node,
     ) -> None:
-        """
-        Convert scalar literal args of targeted_ops in node n of graph_module
+        """Convert scalar literal args of targeted_ops in node n of graph_module
         into attribute get_attr nodes with registered buffers.
         """
         if n.op != "call_function" or n.target not in self.targeted_ops:
@@ -97,8 +96,8 @@ def _convert_scalar_args(
                 graph_module.graph.erase_node(n)
 
     def handle_control_nodes(self, node: Node, graph_module: GraphModule) -> None:
-        """
-        Apply scalar argument conversion on subgraphs of control-flow nodes.
+        """Apply scalar argument conversion on subgraphs of control-flow
+        nodes.
         """
         for _, submodule, _ in get_cond_while_submodules(graph_module):
             for submodule_node in submodule.graph.nodes:

backends/arm/_passes/to_tosa_memory_format_pass.py

@@ -23,8 +23,8 @@
 
 
 def _is_input(node: torch.fx.Node, exported_program: ExportedProgram) -> bool:
-    """
-    Returns True if the node is an input node, i.e. a placeholder or a parameter.
+    """Returns True if the node is an input node, i.e. a placeholder or a
+    parameter.
     """
     return node.op == "placeholder" and not is_param_node(exported_program, node)
 
@@ -42,12 +42,16 @@ def _is_transpose_conv2d_weight(node: torch.fx.Node) -> bool:
 
 
 class ToTosaMemoryFormatPass(ArmPass):
-    """
-    Annotates each node with a tosa_dim_order. tosa_dim_order can be seen as a channels-last dim-order
-    that in most cases will be (0, 2, 3, 1) for nodes with 4D-shapes. The pass also inserts backend.tosa.TRANSPOSE
-    when a transition between 3D and 4D/5D tensors happen.
-    The annotated tosa_dim_order is used to permute the node's shape such that it gives a TOSA-compliant shape.
-    This pass also makes other values aware of spatial dimensions required by future operators by back propogating info as required.
+    """Annotates each node with a tosa_dim_order.
+
+    tosa_dim_order can be seen as a channels-last dim-order that in most cases
+    will be (0, 2, 3, 1) for nodes with 4D-shapes. The pass also inserts
+    backend.tosa.TRANSPOSE when a transition between 3D and 4D/5D tensors
+    happen. The annotated tosa_dim_order is used to permute the node's shape
+    such that it gives a TOSA-compliant shape. This pass also makes other values
+    aware of spatial dimensions required by future operators by back propogating
+    info as required.
+
     """
 
     _passes_required_after: Set[Type[ExportPass]] = set()
@@ -58,8 +62,7 @@ def __init__(self, exported_program: ExportedProgram, *args, **kwargs) -> None:
 
     @staticmethod
     def _channels_last_order(rank: int, spatial_rank: int) -> tuple[int, ...]:
-        """
-        Compute the permutation of tensor dimensions corresponding to a
+        """Compute the permutation of tensor dimensions corresponding to a
         "channels_last"-style memory layout for an arbitrary tensor rank.
 
         In standard PyTorch convention:
@@ -85,6 +88,7 @@ def _channels_last_order(rank: int, spatial_rank: int) -> tuple[int, ...]:
             If `rank <= 2`, the function returns the identity order since there
             are no distinct channel/spatial dimensions.
             In practice only rank 4+ tensors will reach this function as the dim order should be fixed for those.
+
         """
         if rank <= 2:
             return tuple(range(rank))
@@ -96,11 +100,11 @@ def _channels_last_order(rank: int, spatial_rank: int) -> tuple[int, ...]:
 
     @staticmethod
     def _channels_last_inverse_order(rank: int, spatial_rank: int) -> tuple[int, ...]:
-        """
-        Return the inverse permutation of `_channels_last_order`.
+        """Return the inverse permutation of `_channels_last_order`.
+
+        This provides the axis order needed to map a tensor from "channels_last"
+        layout back to its original layout.
 
-        This provides the axis order needed to map a tensor from
-        "channels_last" layout back to its original layout.
         """
         order = ToTosaMemoryFormatPass._channels_last_order(rank, spatial_rank)
         inverse = [0] * rank
@@ -109,15 +113,16 @@ def _channels_last_inverse_order(rank: int, spatial_rank: int) -> tuple[int, ...
         return tuple(inverse)
 
     def _initial_spatial_rank(self, node: torch.fx.Node) -> int:
-        """
-        Infer the initial spatial rank based on the current rank, input node spatial
-        ranks and node target. A spatial dimension includes Height, Width or Depth
-        fields. In most operators this will only ever be Height and Width, but for 3D
-        operators such as conv3d this would contain 3 spatial dims.
+        """Infer the initial spatial rank based on the current rank, input node
+        spatial ranks and node target. A spatial dimension includes Height,
+        Width or Depth fields. In most operators this will only ever be Height
+        and Width, but for 3D operators such as conv3d this would contain 3
+        spatial dims.
 
         Spatial rank is the max of any input node spatial ranks and the number of
         trailing spatial dims we need to preserve (rank - 2, capped at 3). This
         decides which axes must stay channels-last when inserting transposes.
+
         """
         tensor = get_first_fake_tensor(node).data
         # Start by assuming 2D when dealing with rank4+ to account for the base case
@@ -150,9 +155,9 @@ def _initial_spatial_rank(self, node: torch.fx.Node) -> int:
 
     @staticmethod
     def memory_format_differs(shape, spatial_rank):
-        """
-        Determine whether a tensor shape would be laid out differently in
-        channels-first ((N)NCHW) versus channels-last ((N)NHWC) memory format.
+        """Determine whether a tensor shape would be laid out differently in
+        channels-first ((N)NCHW) versus channels-last ((N)NHWC) memory
+        format.
         """
         if len(shape) <= 2 or spatial_rank <= 0:
             return False
@@ -168,8 +173,7 @@ def memory_format_differs(shape, spatial_rank):
     def is_channel_reshape(
         input_shape, output_shape, input_spatial_rank, output_spatial_rank
     ):
-        """
-        Check whether a reshape touches the logical channel or consolidated
+        """Check whether a reshape touches the logical channel or consolidated
         batch dimensions, which would invalidate dim-order annotations.
         """
 
@@ -202,8 +206,7 @@ def get_batch_prod_dim(shape, spatial_rank):
 
     @staticmethod
     def insert_input_transpose(node, input_node, graph_module):
-        """
-        Ensure an input tensor is converted to channels-last ordering by
+        """Ensure an input tensor is converted to channels-last ordering by
         inserting (or folding) a backend `TRANSPOSE` node.
         """
         if input_node.target == exir_ops.backend.tosa.TRANSPOSE.default:
@@ -240,8 +243,7 @@ def insert_input_transpose(node, input_node, graph_module):
 
     @staticmethod
     def insert_output_transpose(node, graph_module):
-        """
-        Convert a producer's output to channels-last by appending a backend
+        """Convert a producer's output to channels-last by appending a backend
         `TRANSPOSE` node and rewiring its users.
         """
 
@@ -280,9 +282,9 @@ def insert_output_transpose(node, graph_module):
     def _insert_view_transpose(
         input_shape, output_shape, node, input_node, graph_module
     ):
-        """
-        Insert the necessary input/output transposes around reshapes that cross
-        the (N)NCHW -> (N)NHWC boundary or that touch channel dimensions.
+        """Insert the necessary input/output transposes around reshapes that
+        cross the (N)NCHW -> (N)NHWC boundary or that touch channel
+        dimensions.
         """
         nchw_to_nhwc = len(input_shape) < 4 and len(output_shape) >= 4
         nhwc_to_nchw = len(input_shape) >= 4 and len(output_shape) < 4
@@ -310,8 +312,10 @@ def _insert_view_transpose(
             ToTosaMemoryFormatPass.insert_output_transpose(node, graph_module)
 
     def insert_tosa_transposes(self, graph_module: torch.fx.GraphModule):
-        """
-        Transposes are needed for operators transforming the input to a different rank, as 4D and 5D-tensors are assumed to be in (N)NHWC-format, whereas all other are in (N)NCHW format.
+        """Transposes are needed for operators transforming the input to a
+        different rank, as 4D and 5D-tensors are assumed to be in (N)NHWC-
+        format, whereas all other are in (N)NCHW format.
+
         This is relevant for the following cases:
         - view:       <4D ->  >=4D
         - view:      >=4D ->   <4D
@@ -321,6 +325,7 @@ def insert_tosa_transposes(self, graph_module: torch.fx.GraphModule):
         - H == W == 1
         - C == 1
         - 1D/2D tensors
+
         """
         for node in graph_module.graph.nodes:
             # call_function and placeholder allowed due to
@@ -383,9 +388,8 @@ def insert_tosa_transposes(self, graph_module: torch.fx.GraphModule):
     def remove_dim_order_kwargs(
         self, graph_module: torch.fx.GraphModule, node: torch.fx.Node
     ):
-        """
-        Drop any user-specified `dim_order` keyword arguments so the pass remains
-        the single source of truth for dim-order annotations.
+        """Drop any user-specified `dim_order` keyword arguments so the pass
+        remains the single source of truth for dim-order annotations.
         """
         if node.op != "call_function":
             return
@@ -439,9 +443,8 @@ def call(self, graph_module: torch.fx.GraphModule):
         return PassResult(graph_module, True)
 
     def _propagate_spatial_ranks(self, nodes):
-        """
-        Propagate `tosa_spatial_rank` metadata backwards so earlier nodes learn
-        about upcoming spatial requirements from future ops.
+        """Propagate `tosa_spatial_rank` metadata backwards so earlier nodes
+        learn about upcoming spatial requirements from future ops.
         """
         changed = True
         while changed:

backends/arm/_passes/unsqueeze_scalar_placeholders_pass.py

@@ -1,4 +1,4 @@
-# Copyright 2024-2025 Arm Limited and/or its affiliates.
+# Copyright 2024-2026 Arm Limited and/or its affiliates.
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
@@ -15,9 +15,11 @@
 
 
 class UnsqueezeScalarPlaceholdersPass(ArmPass):
-    """
-    Placeholders that have node.meta["val"].shape = () cause issues later in the lowering.
+    """Placeholders that have node.meta["val"].shape = () cause issues later in
+    the lowering.
+
     This pass unsqueezes the placeholders to make sure shape is at least (1,).
+
     """
 
     _passes_required_after: Set[Type[ExportPass]] = set()