Cleanup/improve docstrings (#250)

Author: jansel

helion/_compiler/ast_read_writes.py

@@ -62,9 +62,12 @@ def from_ast(node: ast.AST) -> ReadWrites:
         This function traverses the given AST node and collects information
         about variable reads and writes using the `_ReadWriteVisitor` class.
 
-        :param node: The root AST node to analyze.
-        :return: A `ReadWrites` object containing dictionaries of read and
-                 written variable names.
+        Args:
+            node: The root AST node to analyze.
+
+        Returns:
+            A `ReadWrites` object containing dictionaries of read and
+            written variable names.
         """
         visitor = _ReadWriteVisitor()
         visitor.visit(node)
@@ -87,9 +90,12 @@ def ast_rename(node: _A, renames: dict[str, str]) -> _A:
     This function traverses the given AST node and renames variables
     based on the provided mapping of old names to new names.
 
-    :param node: The root AST node to rename variables in.
-    :param renames: A dictionary mapping old variable names to new variable names.
-    :return: The modified AST node with variables renamed.
+    Args:
+        node: The root AST node to rename variables in.
+        renames: A dictionary mapping old variable names to new variable names.
+
+    Returns:
+        The modified AST node with variables renamed.
     """
     visitor = _RenameVisitor(renames)
     visitor.visit(node)
@@ -105,8 +111,11 @@ def visit_Assign(self, node: ast.Assign) -> ast.Assign | None:
         """
         Visit an assignment node and remove it if the target variable is in the to_remove set.
 
-        :param node: The assignment node to visit.
-        :return: The modified assignment node, or None if it should be removed.
+        Args:
+            node: The assignment node to visit.
+
+        Returns:
+            The modified assignment node, or None if it should be removed.
         """
         if len(node.targets) == 1:
             (target,) = node.targets

helion/_compiler/generate_ast.py

@@ -352,9 +352,12 @@ def codegen_precompile_def(
     The precompile function is the same as the normal function, but the call to the
     kernel is replaced with a call to make_precompiler.
 
-    :param host_def: The host function definition to that is used to call the kernel.
-    :param device_function_name: The name of the device function to be called.
-    :return: A transformed function definition with the kernel call replaced.
+    Args:
+        host_def: The host function definition to that is used to call the kernel.
+        device_function_name: The name of the device function to be called.
+
+    Returns:
+        A transformed function definition with the kernel call replaced.
     """
 
     def transform(node: ExtendedAST) -> ExtendedAST:

helion/_compiler/inductor_lowering.py

@@ -257,10 +257,13 @@ def strip_unused_inputs(
     Remove unused inputs from the node.  Inplace updates node.args and
     node.kwargs to replace unused inputs with None.
 
-    :param node: Node to mutate args of
-    :param used_input_names: Set of input names that are used in the node's lowering.
-    :param input_names: Mapping of node inputs to their names.
-    :return:  List of nodes that were used in the lowering.
+    Args:
+        node: Node to mutate args of
+        used_input_names: Set of input names that are used in the node's lowering.
+        input_names: Mapping of node inputs to their names.
+
+    Returns:
+        list[str]: List of names that were used in the lowering.
     """
 
     def mask_unused_inputs(n: torch.fx.Node) -> torch.fx.Node | None:

helion/_compiler/output_header.py

@@ -40,8 +40,11 @@ def get_needed_imports(root: ast.AST) -> str:
     library imports are required based on the variables that are read. It then constructs
     and returns the corresponding import statements.
 
-    :param root: The root AST node to analyze.
-    :return: A string containing the required import statements, separated by newlines.
+    Args:
+        root: The root AST node to analyze.
+
+    Returns:
+        A string containing the required import statements, separated by newlines.
     """
     rw = ReadWrites.from_ast(root)
     result = [library_imports[name] for name in library_imports if name in rw.reads]
@@ -57,8 +60,11 @@ def assert_no_conflicts(fn: FunctionType) -> None:
     not conflict with any reserved names used in the library imports. If
     a conflict is found, an exception is raised.
 
-    :param fn: The function to check for naming conflicts.
-    :raises helion.exc.NamingConflict: If a naming conflict is detected.
+    Args:
+        fn: The function to check for naming conflicts.
+
+    Raises:
+        helion.exc.NamingConflict: If a naming conflict is detected.
     """
     for name in fn.__code__.co_varnames:
         if name in library_imports:
@@ -80,5 +86,8 @@ def assert_no_conflicts(fn: FunctionType) -> None:
 def reserved_names() -> list[str]:
     """
     Retrieve a list of reserved names used in the library imports.
+
+    Returns:
+        A list of reserved names used in the library imports.
     """
     return [*library_imports]

helion/_compiler/variable_origin.py

@@ -29,63 +29,71 @@ def is_global(self) -> bool:
         """
         Check if the origin is a global variable.
 
-        :return: True if the origin is from a global variable, False otherwise.
+        Returns:
+            bool: True if the origin is from a global variable, False otherwise.
         """
         return issubclass(self.base_type(), GlobalOrigin)
 
     def is_argument(self) -> bool:
         """
         Check if the origin is an argument.
 
-        :return: True if the origin is from an argument, False otherwise.
+        Returns:
+            bool: True if the origin is from an argument, False otherwise.
         """
         return issubclass(self.base_type(), ArgumentOrigin)
 
     def is_device(self) -> bool:
         """
         Check if the origin is a device.
 
-        :return: True if the origin is a device, False otherwise.
+        Returns:
+            bool: True if the origin is a device, False otherwise.
         """
         return not self.is_host()
 
     def base_type(self) -> type[Origin]:
         """
         Get the base type of the origin, unwrapping things like attributes.
 
-        :return: The base type of the origin.
+        Returns:
+            type[Origin]: The base type of the origin.
         """
         return type(self)
 
     def needs_rename(self) -> bool:
         """
         Check if the origin needs to be renamed (globals and closures).
 
-        :return: True if the origin needs to be renamed, False otherwise.
+        Returns:
+            bool: True if the origin needs to be renamed, False otherwise.
         """
         return self.is_global()
 
     def depth(self) -> int:
         """
         Get the depth of the origin.
 
-        :return: The depth of the origin, which is 1 by default and increases each wrapper.
+        Returns:
+            int: The depth of the origin, which is 1 by default and increases each wrapper.
         """
         return 1
 
     def host_str(self) -> str:
         """
         Get a string representation of the host origin.
 
-        :raises NotImplementedError: Always raises this error as it should be implemented by subclasses.
+        Raises:
+            NotImplementedError: Always raises this error as it should be implemented by subclasses.
         """
         raise NotImplementedError(type(self).__name__)
 
     def suggest_var_name(self) -> str:
         """
         Suggest a variable name based on the origin.
 
-        :raises NotImplementedError: Always raises this error as it should be implemented by subclasses.
+        Raises:
+            NotImplementedError: Always raises this error as it should be implemented by subclasses.
         """
         raise NotImplementedError(type(self).__name__)