Arm backend: Fixup and make docstrings consistent for key functions (pytorch#15443)

* Align formatting across key functions and classes of backend
* Fix some docstrings

Signed-off-by: Tom Allsop <[email protected]>

Author: tom-arm

backends/arm/common/arm_compile_spec.py

@@ -173,18 +173,30 @@ def to_list(self):
         return compile_spec
 
     def get_intermediate_path(self) -> str | None:
+        """
+        Gets the path used for dumping intermediate results such as tosa and pte.
+
+        Returns:
+            Path where intermediate results are saved.
+        """
         return self.path_for_intermediates
 
     def dump_intermediate_artifacts_to(self, output_path: str | None):
         """
         Sets a path for dumping intermediate results during such as tosa and pte.
+
+        Args:
+            output_path: Path to dump intermediate results to.
         """
         self.path_for_intermediates = output_path
         return self
 
     def dump_debug_info(self, debug_mode: DebugMode | None):
         """
-        Dump debugging information into the intermediates path
+        Dump debugging information into the intermediates path.
+
+        Args:
+            debug_mode: The debug mode to use for dumping debug information.
         """
         self.tosa_debug_mode = debug_mode
         return self

backends/arm/ethosu/compile_spec.py

@@ -15,6 +15,16 @@
 
 
 class EthosUCompileSpec(ArmCompileSpec):
+    """
+    Compile spec for Ethos-U NPU.
+
+    Args:
+        target: Ethos-U accelerator configuration, e.g. ethos-u55-128.
+        system_config: System configuration to select from the Vela configuration file.
+        memory_mode: Memory mode to select from the Vela configuration file.
+        extra_flags: Extra flags for the Vela compiler.
+        config_ini: Vela configuration file(s) in Python ConfigParser .ini file format.
+    """
 
     _TARGET_KEY = "target"
 
@@ -26,16 +36,6 @@ def __init__(
         extra_flags: list[str] | None = None,
         config_ini: str | None = "Arm/vela.ini",
     ):
-        """Generate compile spec for Ethos-U NPU
-        Args:
-        target: Ethos-U accelerator configuration, e.g. ethos-u55-128
-        system_config: System configuration to select from the Vela
-            configuration file
-        memory_mode: Memory mode to select from the Vela configuration file
-        extra_flags: Extra flags for the Vela compiler
-        config_ini: Vela configuration file(s) in Python ConfigParser .ini
-            file format
-        """
         self.target = target
 
         # Set vela compiler flags
@@ -87,6 +87,7 @@ def from_list_hook(cls, compile_spec, specs: dict[str, str]):
         compile_spec.target = specs.get(cls._TARGET_KEY, None)
 
     def validate(self):
+        """Throws an error if the compile spec is not valid."""
         if len(self.compiler_flags) == 0:
             raise ValueError(
                 "compile_flags are required in the CompileSpec list for EthosUBackend"

backends/arm/ethosu/partitioner.py

@@ -17,7 +17,7 @@ class EthosUPartitioner(TOSAPartitioner):
     """
     Partitions subgraphs supported by the Arm Ethos-U backend.
 
-    Attributes:
+    Args:
         compile_spec: List of CompileSpec objects for Ethos-U backend.
         additional_checks: Optional sequence of additional operator support checks.
     """

backends/arm/quantizer/arm_quantizer.py

@@ -333,41 +333,67 @@ def __init__(
         self.module_name_config: Dict[str, Optional[QuantizationConfig]] = {}
 
     def set_global(self, quantization_config: QuantizationConfig) -> TOSAQuantizer:
-        """Set quantization_config for submodules that are not already annotated by name or type filters."""
+        """
+        Set quantization_config for submodules that are not already annotated by name or type filters.
+
+        Args:
+            quantization_config: The QuantizationConfig to set as global configuration.
+        """
         self.global_config = quantization_config
         return self
 
     def set_module_type(
         self, module_type: Callable, quantization_config: QuantizationConfig
     ) -> TOSAQuantizer:
-        """Set quantization_config for a submodule with type: `module_type`, for example:
+        """
+        Set quantization_config for a submodule with type: `module_type`, for example:
         quantizer.set_module_name(Sub) or quantizer.set_module_name(nn.Linear), it will quantize all supported operator/operator
-        patterns in the submodule with this module type with the given `quantization_config`
+        patterns in the submodule with this module type with the given `quantization_config`.
+
+        Args:
+            module_type: The type of the submodule to set the quantization config for.
+            quantization_config: The QuantizationConfig to set for the submodule.
         """
         self.module_type_config[module_type] = quantization_config
         return self
 
     def set_module_name(
         self, module_name: str, quantization_config: Optional[QuantizationConfig]
     ) -> TOSAQuantizer:
-        """Set quantization_config for a submodule with name: `module_name`, for example:
+        """
+        Set quantization_config for a submodule with name: `module_name`, for example:
         quantizer.set_module_name("blocks.sub"), it will quantize all supported operator/operator
         patterns in the submodule with this module name with the given `quantization_config`
+
+        Args:
+            module_name: The name of the submodule to set the quantization config for.
+            quantization_config: The QuantizationConfig to set for the submodule.
         """
         # Validate that quantization_config is provided
         if quantization_config is None:
             raise ValueError("quantization_config == None is not supported yet")
         self.module_name_config[module_name] = quantization_config
         return self
 
-    def set_io(self, quantization_config):
-        """Set quantization_config for input and output nodes."""
+    def set_io(self, quantization_config: QuantizationConfig) -> TOSAQuantizer:
+        """
+        Set quantization_config for input and output nodes.
+
+        Args:
+            quantization_config: The QuantizationConfig to set for input and output nodes.
+        """
         self.io_config = quantization_config
         return self
 
     def transform_for_annotation(self, model: GraphModule) -> GraphModule:
-        """An initial pass for transforming the graph to prepare it for annotation.
+        """
+        An initial pass for transforming the graph to prepare it for annotation.
         Currently transforms scalar values to tensor attributes.
+
+        Args:
+            model: The model to transform.
+        Returns:
+            The transformed model.
         """
 
         # TODO: Fix the need to lazily import this.
@@ -465,10 +491,24 @@ def validate(self, model: GraphModule) -> None:
 
 
 class EthosUQuantizer(TOSAQuantizer):
+    """
+    Quantizer supported by the Arm Ethos-U backend.
+
+    Args:
+        compile_spec: A EthosUCompileSpec instance.
+    """
+
     def __init__(self, compile_spec: EthosUCompileSpec) -> None:
         super().__init__(compile_spec)
 
 
 class VgfQuantizer(TOSAQuantizer):
+    """
+    Quantizer supported by the Arm Vgf backend.
+
+    Args:
+        compile_spec: A VgfCompileSpec instance.
+    """
+
     def __init__(self, compile_spec: VgfCompileSpec) -> None:
         super().__init__(compile_spec)

backends/arm/tosa/partitioner.py

@@ -193,7 +193,7 @@ def _tag_module(  # noqa
         """Tag nodes in a module, possibly a submodule, from the containing program.
 
         Args:
-            module: a GraphModule from `containing_program` to tag nodes in.
+            module: A GraphModule from `containing_program` to tag nodes in.
             containing_program: The ExportedProgram that contains the module.
             reporter: A reporter to report why nodes were rejected.
         Returns:

backends/arm/vgf/compile_spec.py

@@ -15,19 +15,19 @@
 
 
 class VgfCompileSpec(ArmCompileSpec):
+    """
+    Compile spec for VGF compatible targets.
+
+    Args:
+        tosa_spec: TOSA specification that should be targeted.
+        compiler_flags: Extra compiler flags for converter_backend.
+    """
 
     def __init__(
         self,
         tosa_spec: TosaSpecification | str | None = None,
         compiler_flags: list[str] | None = None,
     ):
-        """
-        Generate compile spec for VGF compatible targets
-
-        Args:
-            compiler_flags: Extra compiler flags for converter_backend
-        """
-
         if tosa_spec is None:
             tosa_spec = "TOSA-1.0+FP"
         if isinstance(tosa_spec, str):

backends/arm/vgf/partitioner.py

@@ -14,6 +14,14 @@
 
 @final
 class VgfPartitioner(TOSAPartitioner):
+    """
+    Partitions subgraphs supported by the Arm Vgf backend.
+
+    Args:
+        compile_spec: The Vgf compilation specification.
+        additional_checks: Optional sequence of additional operator support checks.
+    """
+
     def __init__(
         self,
         compile_spec: VgfCompileSpec,