Arm backend: Format docs in backends/arm/test/passes (#17376)

Sebastian-Larsson · web-flow · commit ccb13dde1afe · 2026-02-12T09:48:10.000+01:00
Signed-off-by: Sebastian Larsson &lt;sebastian.larsson@arm.com&gt;
diff --git a/backends/arm/test/passes/test_convert_expand_copy_to_repeat.py b/backends/arm/test/passes/test_convert_expand_copy_to_repeat.py
@@ -1,4 +1,4 @@
-# Copyright 2025 Arm Limited and/or its affiliates.
+# Copyright 2025-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.
@@ -16,9 +16,7 @@
 
 
 class Expand(torch.nn.Module):
-    """
-    Basic expand model using torch.Tensor.expand function
-    """
+    """Basic expand model using torch.Tensor.expand function."""
 
     def __init__(self) -> None:
         super().__init__()
diff --git a/backends/arm/test/passes/test_convert_int64_const_ops_to_int32.py b/backends/arm/test/passes/test_convert_int64_const_ops_to_int32.py
@@ -1,4 +1,4 @@
-# Copyright 2025 Arm Limited and/or its affiliates.
+# Copyright 2025-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.
@@ -361,11 +361,13 @@ def test_convert_int64_const_ops_to_int32_tosa_FP_full(
 def test_convert_int64_const_ops_to_int32_tosa_INT_full(
     test_data: input_t2,
 ) -> None:
-    """
-    For INT profile, _lifted_tensor_constant0 is still int64 after applying ConvertInt64ConstOpsToInt32Pass().
+    """For INT profile, _lifted_tensor_constant0 is still int64 after applying
+    ConvertInt64ConstOpsToInt32Pass().
+
     And an int64->int32 cast is inserted at the beginning of the graph.
     TODO: Explore why _lifted_tensor_constant0 is handled in different ways in FP and INT profile.
         Find a way to optimize out the int64->int32 cast.
+
     """
     module = FullIncrementViewMulXLessThanY()
     aten_ops_checks = [
diff --git a/backends/arm/test/passes/test_convert_split_to_slice.py b/backends/arm/test/passes/test_convert_split_to_slice.py
@@ -1,4 +1,4 @@
-# Copyright 2025 Arm Limited and/or its affiliates.
+# Copyright 2025-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.
@@ -22,9 +22,7 @@ def get_inputs(self) -> input_t: ...
 
 
 class Split(torch.nn.Module):
-    """
-    Basic split model using torch.split function
-    """
+    """Basic split model using torch.split function."""
 
     def get_inputs(self) -> input_t:
         return (torch.rand(10),)
@@ -34,9 +32,7 @@ def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, ...]:
 
 
 class SplitTensor(torch.nn.Module):
-    """
-    Basic split model using torch.Tensor.split function
-    """
+    """Basic split model using torch.Tensor.split function."""
 
     def get_inputs(self) -> input_t:
         return (torch.rand(10),)
diff --git a/backends/arm/test/passes/test_decompose_avg_pool2d_pass.py b/backends/arm/test/passes/test_decompose_avg_pool2d_pass.py
@@ -1,4 +1,4 @@
-# Copyright 2025 Arm Limited and/or its affiliates.
+# Copyright 2025-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.
@@ -20,9 +20,7 @@ def get_inputs(self) -> input_t: ...
 
 
 class AvgPool2dWithStride(torch.nn.Module):
-    """
-    avg_pool2d model with explicit stride parameter
-    """
+    """avg_pool2d model with explicit stride parameter."""
 
     def get_inputs(self) -> input_t:
         return (torch.rand(1, 3, 8, 8),)
@@ -32,8 +30,8 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 
 
 class AvgPool2dWithoutStride(torch.nn.Module):
-    """
-    avg_pool2d model without stride parameter (should default to kernel_size)
+    """avg_pool2d model without stride parameter (should default to
+    kernel_size)
     """
 
     def get_inputs(self) -> input_t:
@@ -44,9 +42,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 
 
 class AvgPool2dListKernel(torch.nn.Module):
-    """
-    avg_pool2d model with list kernel_size and no stride
-    """
+    """avg_pool2d model with list kernel_size and no stride."""
 
     def get_inputs(self) -> input_t:
         return (torch.rand(1, 3, 8, 8),)
@@ -64,7 +60,9 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 
 @common.parametrize("module", modules)
 def test_decompose_avg_pool2d_tosa_FP(module: ModuleWithInputs) -> None:
-    """Test that DecomposeAvgPool2d pass works correctly with and without stride parameters."""
+    """Test that DecomposeAvgPool2d pass works correctly with and without stride
+    parameters.
+    """
     nn_module = cast(torch.nn.Module, module)
     pipeline = PassPipeline[input_t](
         nn_module,
diff --git a/backends/arm/test/passes/test_decompose_div_pass.py b/backends/arm/test/passes/test_decompose_div_pass.py
@@ -1,4 +1,4 @@
-# Copyright 2025 Arm Limited and/or its affiliates.
+# Copyright 2025-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.
@@ -20,9 +20,7 @@ def get_inputs(self) -> input_t: ...
 
 
 class Div(torch.nn.Module):
-    """
-    Basic div model using torch.div
-    """
+    """Basic div model using torch.div."""
 
     def get_inputs(self) -> input_t:
         return (torch.rand(10),)
@@ -32,9 +30,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 
 
 class DivTensor(torch.nn.Module):
-    """
-    Basic div model using torch.Tensor.div
-    """
+    """Basic div model using torch.Tensor.div."""
 
     def get_inputs(self) -> input_t:
         return (torch.rand(10),)
diff --git a/backends/arm/test/passes/test_decompose_int_pow_pass.py b/backends/arm/test/passes/test_decompose_int_pow_pass.py
@@ -1,4 +1,4 @@
-# Copyright 2025 Arm Limited and/or its affiliates.
+# Copyright 2025-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.
@@ -23,9 +23,7 @@ def get_inputs(self) -> input_t: ...
 
 
 class Square(torch.nn.Module):
-    """
-    Basic squaring
-    """
+    """Basic squaring."""
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         return x.square()
@@ -35,9 +33,7 @@ def get_inputs(self) -> input_t:
 
 
 class Pow(torch.nn.Module):
-    """
-    Basic squaring
-    """
+    """Basic squaring."""
 
     def __init__(self, exponent: int) -> None:
         super().__init__()
diff --git a/backends/arm/test/passes/test_decompose_layernorm_pass.py b/backends/arm/test/passes/test_decompose_layernorm_pass.py
@@ -1,4 +1,4 @@
-# Copyright 2025 Arm Limited and/or its affiliates.
+# Copyright 2025-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.
@@ -16,9 +16,7 @@
 
 
 class LayerNorm(torch.nn.Module):
-    """
-    Basic layer_norm model using torch.nn.layer_norm layer
-    """
+    """Basic layer_norm model using torch.nn.layer_norm layer."""
 
     def __init__(self):
         super(LayerNorm, self).__init__()
diff --git a/backends/arm/test/passes/test_decompose_linalg_vector_norm_pass.py b/backends/arm/test/passes/test_decompose_linalg_vector_norm_pass.py
@@ -1,4 +1,4 @@
-# Copyright 2025 Arm Limited and/or its affiliates.
+# Copyright 2025-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.
@@ -64,12 +64,14 @@ def get_inputs(self) -> input_t:
 
 @common.parametrize("module", modules)
 def test_decompose_linalg_vector_norm_tosa_INT(module: ModuleWithInputs) -> None:
-    """
-    This test creates a PassPipeline that applies the DecomposeLinalgVectorNormPass.
+    """This test creates a PassPipeline that applies the
+    DecomposeLinalgVectorNormPass.
+
     The expected primitive ops vary depending on the norm order:
       - p == 1: should decompose to ABS and SUM.
       - p == 2 (default): should decompose to MUL, SUM, and SQRT.
       - Other p: should decompose to ABS, two instances of POW, and SUM.
+
     """
     ord_val = module.ord if module.ord is not None else 2.0
 
diff --git a/backends/arm/test/passes/test_decompose_meandim_pass.py b/backends/arm/test/passes/test_decompose_meandim_pass.py
@@ -1,4 +1,4 @@
-# Copyright 2025 Arm Limited and/or its affiliates.
+# Copyright 2025-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.
@@ -27,9 +27,7 @@ def get_inputs(self) -> input_t: ...
 
 
 class MeanDim(torch.nn.Module):
-    """
-    Basic mean model using torch.mean with keepdim = True
-    """
+    """Basic mean model using torch.mean with keepdim = True."""
 
     ops_after_pass = u55_ops_after_pass = {
         "torch.ops.aten.sum.dim_IntList": 2,
@@ -53,9 +51,7 @@ def get_inputs(self) -> input_t:
 
 
 class MeanDimTensor(torch.nn.Module):
-    """
-    Basic mean model using torch.Tensor.mean with keepdim = False
-    """
+    """Basic mean model using torch.Tensor.mean with keepdim = False."""
 
     ops_after_pass = {
         "torch.ops.aten.sum.dim_IntList": 2,
diff --git a/backends/arm/test/passes/test_decompose_softmax_pass.py b/backends/arm/test/passes/test_decompose_softmax_pass.py
@@ -1,4 +1,4 @@
-# Copyright 2025 Arm Limited and/or its affiliates.
+# Copyright 2025-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.
@@ -14,9 +14,7 @@
 
 
 class Softmax(torch.nn.Module):
-    """
-    Basic torch.nn.softmax layer model
-    """
+    """Basic torch.nn.softmax layer model."""
 
     def __init__(self):
         super(Softmax, self).__init__()
@@ -31,9 +29,7 @@ def get_inputs(self) -> input_t:
 
 
 class SoftmaxLog(torch.nn.Module):
-    """
-    Basic torch.nn.log_softmax layer model
-    """
+    """Basic torch.nn.log_softmax layer model."""
 
     def __init__(self):
         super(SoftmaxLog, self).__init__()
diff --git a/backends/arm/test/passes/test_decompose_var_pass.py b/backends/arm/test/passes/test_decompose_var_pass.py
@@ -1,4 +1,4 @@
-# Copyright 2025 Arm Limited and/or its affiliates.
+# Copyright 2025-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.
@@ -20,9 +20,7 @@ def get_inputs(self) -> input_t: ...
 
 
 class VarDim(torch.nn.Module):
-    """
-    Basic variance model using torch.Tensor.var function.
-    """
+    """Basic variance model using torch.Tensor.var function."""
 
     def __init__(self, keepdim):
         super(VarDim, self).__init__()
@@ -36,9 +34,7 @@ def get_inputs(self) -> input_t:
 
 
 class VarCorrection(torch.nn.Module):
-    """
-    Basic variance model using torch.var function.
-    """
+    """Basic variance model using torch.var function."""
 
     def __init__(self, keepdim):
         super(VarCorrection, self).__init__()
diff --git a/backends/arm/test/passes/test_decorate_fp32_to_int32_casting_pass.py b/backends/arm/test/passes/test_decorate_fp32_to_int32_casting_pass.py
@@ -1,4 +1,4 @@
-# Copyright 2025 Arm Limited and/or its affiliates.
+# Copyright 2025-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.
@@ -66,10 +66,12 @@ def test_decorate_fp32_to_int32_casting_tosa_FP(
 def test_decorate_fp32_to_int32_casting_tosa_INT(
     test_data: Callable[[], Tuple[torch.Tensor, torch.dtype]]
 ) -> None:
-    """
-    Casting operation involving floating-point dtypes will be rejected in INT/INT profile.
-    Therefore, the DecorateFp32toInt32CastingPass is not required in this profile.
-    Add a INT test to ensure that such casting is rejected as expected.
+    """Casting operation involving floating-point dtypes will be rejected in
+    INT/INT profile.
+
+    Therefore, the DecorateFp32toInt32CastingPass is not required in this
+    profile. Add a INT test to ensure that such casting is rejected as expected.
+
     """
     test_tensor, target_dtype = test_data()
     module = FP32ToINT32Casting(target_dtype)
diff --git a/backends/arm/test/passes/test_fold_qdq_pass.py b/backends/arm/test/passes/test_fold_qdq_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.
@@ -25,12 +25,12 @@ def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
 
 @common.parametrize("test_data", SimpleQuantizeModel.test_data)
 def test_fold_and_annotate_q_params_tosa_INT(test_data: input_t) -> None:
-    """
-    Tests the FoldAndAnnotateQParamsPass which folds dq/q nodes into
-    the node and stores the quantization parameters in meta.
+    """Tests the FoldAndAnnotateQParamsPass which folds dq/q nodes into the node
+    and stores the quantization parameters in meta.
+
+    Check that the pass runs for add operation and that one q node and one dq
+    node is removed from the representation.
 
-    Check that the pass runs for add operation and that one q node and one dq node
-    is removed from the representation.
     """
     module = SimpleQuantizeModel()
     pipeline = PassPipeline[input_t](
diff --git a/backends/arm/test/passes/test_fuse_batchnorm_pass.py b/backends/arm/test/passes/test_fuse_batchnorm_pass.py
@@ -1,4 +1,4 @@
-# Copyright 2025 Arm Limited and/or its affiliates.
+# Copyright 2025-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.
@@ -150,8 +150,9 @@ def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
 
 @common.parametrize("module", modules)
 def test_fuse_batch_norm2d_tosa_FP(module: ModuleWithBatchNormAttrs) -> None:
-    """Test various cases where the batchnorm should either be fused with a previous
-    conv, or converted to a new conv."""
+    """Test various cases where the batchnorm should either be fused with a
+    previous conv, or converted to a new conv.
+    """
     nn_module = cast(torch.nn.Module, module)
     pipeline = PassPipeline[input_t](
         nn_module,
diff --git a/backends/arm/test/passes/test_ioquantization_pass.py b/backends/arm/test/passes/test_ioquantization_pass.py
@@ -28,8 +28,8 @@ def forward(self, x, y):
 
 @common.parametrize("test_data", SimpleModel.test_data)
 def test_quantize_io_u55_INT(test_data: input_t):
-    """
-    Test the executorch/exir/passes/quantize_io_pass pass works(meaning we don't get Q/DQ nodes) on a simple model
+    """Test the executorch/exir/passes/quantize_io_pass pass works(meaning we
+    don't get Q/DQ nodes) on a simple model.
     """
     model = SimpleModel()
     pipeline = EthosU55PipelineINT(
diff --git a/backends/arm/test/passes/test_remove_clone_pass.py b/backends/arm/test/passes/test_remove_clone_pass.py
@@ -1,4 +1,4 @@
-# Copyright 2025 Arm Limited and/or its affiliates.
+# Copyright 2025-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.
@@ -14,9 +14,7 @@
 
 
 class Clone(torch.nn.Module):
-    """
-    Basic remove layer model to test RemoveNoopePass
-    """
+    """Basic remove layer model to test RemoveNoopePass."""
 
     def __init__(self):
         super(Clone, self).__init__()
diff --git a/backends/arm/test/passes/test_replace_inf_values_pass.py b/backends/arm/test/passes/test_replace_inf_values_pass.py
diff --git a/backends/arm/test/passes/test_rescale_pass.py b/backends/arm/test/passes/test_rescale_pass.py
diff --git a/backends/arm/test/passes/test_to_tosa_memory_format.py b/backends/arm/test/passes/test_to_tosa_memory_format.py
diff --git a/backends/arm/test/passes/test_unsqueeze_before_repeat_pass.py b/backends/arm/test/passes/test_unsqueeze_before_repeat_pass.py
