Apply formatting

adam-smnk · adam-smnk · commit e8f80fa34220 · 2025-11-24T13:00:18.000+01:00
diff --git a/ingress/mlir-gen/mlir_gen/__main__.py b/ingress/mlir-gen/mlir_gen/__main__.py
@@ -1,4 +1,5 @@
 import sys
+
 from . import main
 
 # Invoke on command line with `python -m mlir_gen`.
diff --git a/ingress/mlir-gen/mlir_gen/einsum.py b/ingress/mlir-gen/mlir_gen/einsum.py
@@ -1,10 +1,16 @@
 from typing import Union
 
 from mlir import ir
-from mlir.dialects import arith, linalg, tensor
-
-from . import named, generic
-from .utils import get_outputs, get_weights, get_bias, affine_map
+from mlir.dialects import arith
+from mlir.dialects import linalg
+from mlir.dialects import tensor
+
+from . import generic
+from . import named
+from .utils import affine_map
+from .utils import get_bias
+from .utils import get_outputs
+from .utils import get_weights
 
 
 def times_weights(
diff --git a/ingress/mlir-gen/mlir_gen/generic.py b/ingress/mlir-gen/mlir_gen/generic.py
@@ -1,16 +1,17 @@
 from typing import Union
 
 from mlir import ir
-from mlir.dialects import linalg, arith, tensor, math
-
-from .utils import (
-    affine_map,
-    get_bias,
-    get_outputs,
-    get_weights,
-    parallel,
-    reduction,
-)
+from mlir.dialects import arith
+from mlir.dialects import linalg
+from mlir.dialects import math
+from mlir.dialects import tensor
+
+from .utils import affine_map
+from .utils import get_bias
+from .utils import get_outputs
+from .utils import get_weights
+from .utils import parallel
+from .utils import reduction
 
 
 def affine_maps_and_iter_types(rank: int):
diff --git a/ingress/mlir-gen/mlir_gen/main.py b/ingress/mlir-gen/mlir_gen/main.py
@@ -1,17 +1,21 @@
-import random
-import sys
-
 from argparse import ArgumentParser
-from typing import Sequence, Dict, Any, Optional
 from collections import namedtuple
+import random
+import sys
+from typing import Any
+from typing import Dict
+from typing import Optional
+from typing import Sequence
 
 import numpy as np
 
 from mlir import ir
 from mlir.dialects import func
 
-from . import named, generic, einsum, utils as gen_utils
-
+from . import einsum
+from . import generic
+from . import named
+from . import utils as gen_utils
 
 BlockFactors = namedtuple("BlockFactors", "m n k vnni")
 
@@ -136,9 +140,9 @@ def weights(
             assert k_as_num_inputs % block.k == 0, "invalid tile size for K dim"
             assert n_as_num_outputs % block.n == 0, "invalid tile size for N dim"
             if block.vnni:
-                assert (
-                    block.n % block.vnni == 0
-                ), "incompatible tile sizes for N and VNNI dims"
+                assert block.n % block.vnni == 0, (
+                    "incompatible tile sizes for N and VNNI dims"
+                )
                 shape = (
                     n_as_num_outputs // block.n,
                     k_as_num_inputs // block.k,
diff --git a/ingress/mlir-gen/mlir_gen/named.py b/ingress/mlir-gen/mlir_gen/named.py
@@ -1,9 +1,13 @@
 from typing import Union
 
 from mlir import ir
-from mlir.dialects import linalg, tensor, arith
+from mlir.dialects import arith
+from mlir.dialects import linalg
+from mlir.dialects import tensor
 
-from .utils import get_outputs, get_weights, get_bias
+from .utils import get_bias
+from .utils import get_outputs
+from .utils import get_weights
 
 
 def times_weights(
diff --git a/ingress/mlir-gen/mlir_gen/utils.py b/ingress/mlir-gen/mlir_gen/utils.py
@@ -1,12 +1,15 @@
-import struct
-from enum import Enum, auto
 from collections import abc
+from enum import Enum
+from enum import auto
+import struct
 from typing import Union
 
 import numpy as np
 
 from mlir import ir
-from mlir.dialects import arith, linalg, tensor
+from mlir.dialects import arith
+from mlir.dialects import linalg
+from mlir.dialects import tensor
 
 
 class ConstantInitKind(Enum):
diff --git a/python/examples/ingress/torch/MLPModel/model.py b/python/examples/ingress/torch/MLPModel/model.py
@@ -3,16 +3,11 @@
 import torch
 import torch.nn as nn
 
-import os
 
 class MLPModel(nn.Module):
     def __init__(self):
         super().__init__()
-        self.net = nn.Sequential(
-            nn.Linear(10, 32),
-            nn.ReLU(),
-            nn.Linear(32, 2)
-        )
+        self.net = nn.Sequential(nn.Linear(10, 32), nn.ReLU(), nn.Linear(32, 2))
 
     def forward(self, x):
         return self.net(x)
diff --git a/python/examples/ingress/torch/mlp_from_file.py b/python/examples/ingress/torch/mlp_from_file.py
@@ -17,8 +17,8 @@
 from pathlib import Path
 
 # MLIR infrastructure imports (only needed if you want to manipulate the MLIR module)
-import mlir.dialects.func as func
 from mlir import ir
+import mlir.dialects.func as func
 
 # Lighthouse imports
 from lighthouse.ingress.torch import import_from_file
@@ -34,6 +34,7 @@
 # - Loads the MLPModel class and instantiates it with arguments obtained from 'get_init_inputs()'
 # - Calls get_sample_inputs() to get sample input tensors for shape inference
 # - Converts PyTorch model to linalg-on-tensors dialect operations using torch_mlir
+# fmt: off
 mlir_module_ir: ir.Module = import_from_file(
     model_path,                              # Path to the Python file containing the model
     model_class_name="MLPModel",             # Name of the PyTorch nn.Module class to convert
@@ -42,6 +43,7 @@
     dialect="linalg-on-tensors",             # Target MLIR dialect (linalg ops on tensor types)
     ir_context=ir_context                    # MLIR context for the conversion
 )
+# fmt: on
 
 # The PyTorch model is now converted to MLIR at this point. You can now convert
 # the MLIR module to a text form (e.g. 'str(mlir_module_ir)') and save it to a file.
diff --git a/python/examples/ingress/torch/mlp_from_model.py b/python/examples/ingress/torch/mlp_from_model.py
@@ -15,8 +15,8 @@
 import torch
 
 # MLIR infrastructure imports (only needed if you want to manipulate the MLIR module)
-import mlir.dialects.func as func
 from mlir import ir
+import mlir.dialects.func as func
 
 # Lighthouse imports
 from lighthouse.ingress.torch import import_from_model
@@ -31,9 +31,7 @@
 ir_context = ir.Context()
 # Step 2: Convert the PyTorch model to MLIR
 mlir_module_ir: ir.Module = import_from_model(
-    model,
-    sample_args=(sample_input,),
-    ir_context=ir_context
+    model, sample_args=(sample_input,), ir_context=ir_context
 )
 
 # The PyTorch model is now converted to MLIR at this point. You can now convert
diff --git a/python/examples/mlir/compile_and_run.py b/python/examples/mlir/compile_and_run.py
@@ -1,8 +1,9 @@
 # RUN: %PYTHON %s
 
-import torch
 import argparse
 
+import torch
+
 from mlir import ir
 from mlir.dialects import transform
 from mlir.dialects.transform import structured
diff --git a/python/examples/schedule/transform_a_payload_according_to_a_schedule.py b/python/examples/schedule/transform_a_payload_according_to_a_schedule.py
@@ -3,10 +3,20 @@
 # Simply demonstrates applying a schedule to a payload.
 # To do so generates a basic payload and a basic schedule, purely as an example.
 
-from mlir.ir import Context, Location, InsertionPoint, Operation, Module
-from mlir.ir import RankedTensorType, F32Type, UnitAttr
-from mlir.dialects import arith, func, linalg, tensor, transform
+from mlir.dialects import arith
+from mlir.dialects import func
+from mlir.dialects import linalg
+from mlir.dialects import tensor
+from mlir.dialects import transform
 from mlir.dialects.transform import structured
+from mlir.ir import Context
+from mlir.ir import F32Type
+from mlir.ir import InsertionPoint
+from mlir.ir import Location
+from mlir.ir import Module
+from mlir.ir import Operation
+from mlir.ir import RankedTensorType
+from mlir.ir import UnitAttr
 
 
 def example_payload() -> Module:
diff --git a/python/lighthouse/ingress/torch/__init__.py b/python/lighthouse/ingress/torch/__init__.py
@@ -1,3 +1,9 @@
 """Provides functions to convert PyTorch models to MLIR."""
 
-from .importer import import_from_file, import_from_model
+from .importer import import_from_file
+from .importer import import_from_model
+
+__all__ = [
+    "import_from_file",
+    "import_from_model",
+]
diff --git a/python/lighthouse/ingress/torch/importer.py b/python/lighthouse/ingress/torch/importer.py
@@ -1,9 +1,11 @@
 import importlib
 import importlib.util
 from pathlib import Path
-from typing import Iterable, Mapping
+from typing import Iterable
+from typing import Mapping
 
-from lighthouse.ingress.torch.utils import load_and_run_callable, maybe_load_and_run_callable
+from lighthouse.ingress.torch.utils import load_and_run_callable
+from lighthouse.ingress.torch.utils import maybe_load_and_run_callable
 
 try:
     import torch
@@ -25,6 +27,7 @@
 
 from mlir import ir
 
+
 def import_from_model(
     model: nn.Module,
     sample_args: Iterable,
@@ -49,10 +52,10 @@ def import_from_model(
         ir_context (ir.Context, optional): An optional MLIR context to use for parsing
             the module. If not provided, the module is returned as a string.
         **kwargs: Additional keyword arguments passed to the ``torch_mlir.fx.export_and_import`` function.
-    
+
     Returns:
         str | ir.Module: The imported MLIR module as a string or an ir.Module if `ir_context` is provided.
-    
+
     Examples:
         >>> import torch
         >>> import torch.nn as nn
@@ -61,17 +64,22 @@ def import_from_model(
         ...     def __init__(self):
         ...         super().__init__()
         ...         self.fc = nn.Linear(10, 5)
+        ...
         ...     def forward(self, x):
         ...         return self.fc(x)
         >>> model = SimpleModel()
         >>> sample_input = (torch.randn(1, 10),)
         >>> #
         >>> # option 1: get MLIR module as a string
-        >>> mlir_module : str = import_from_model(model, sample_input, dialect="linalg-on-tensors")
-        >>> print(mlir_module) # prints the MLIR module in linalg-on-tensors dialect
+        >>> mlir_module: str = import_from_model(
+        ...     model, sample_input, dialect="linalg-on-tensors"
+        ... )
+        >>> print(mlir_module)  # prints the MLIR module in linalg-on-tensors dialect
         >>> # option 2: get MLIR module as an ir.Module
         >>> ir_context = ir.Context()
-        >>> mlir_module_ir : ir.Module = import_from_model(model, sample_input, dialect="tosa", ir_context=ir_context)
+        >>> mlir_module_ir: ir.Module = import_from_model(
+        ...     model, sample_input, dialect="tosa", ir_context=ir_context
+        ... )
     """
     if dialect == "linalg":
         raise ValueError(
@@ -134,45 +142,48 @@ def import_from_file(
         ir_context (ir.Context, optional): An optional MLIR context to use for parsing
             the module. If not provided, the module is returned as a string.
         **kwargs: Additional keyword arguments passed to the ``torch_mlir.fx.export_and_import`` function.
-    
+
     Returns:
         str | ir.Module: The imported MLIR module as a string or an ir.Module if `ir_context` is provided.
-    
+
     Examples:
         Given a file `path/to/model_file.py` with the following content:
         ```python
         import torch
         import torch.nn as nn
 
+
         class MyModel(nn.Module):
             def __init__(self):
                 super().__init__()
                 self.fc = nn.Linear(10, 5)
+
             def forward(self, x):
                 return self.fc(x)
 
+
         def get_inputs():
             return (torch.randn(1, 10),)
         ```
 
         The import script would look like:
         >>> from lighthouse.ingress.torch_import import import_from_file
         >>> # option 1: get MLIR module as a string
-        >>> mlir_module : str = import_from_file(
+        >>> mlir_module: str = import_from_file(
         ...     "path/to/model_file.py",
         ...     model_class_name="MyModel",
         ...     init_args_fn_name=None,
-        ...     dialect="linalg-on-tensors"
+        ...     dialect="linalg-on-tensors",
         ... )
-        >>> print(mlir_module) # prints the MLIR module in linalg-on-tensors dialect
+        >>> print(mlir_module)  # prints the MLIR module in linalg-on-tensors dialect
         >>> # option 2: get MLIR module as an ir.Module
         >>> ir_context = ir.Context()
-        >>> mlir_module_ir : ir.Module = import_from_file(
+        >>> mlir_module_ir: ir.Module = import_from_file(
         ...     "path/to/model_file.py",
         ...     model_class_name="MyModel",
         ...     init_args_fn_name=None,
         ...     dialect="linalg-on-tensors",
-        ...     ir_context=ir_context
+        ...     ir_context=ir_context,
         ... )
     """
     if isinstance(filepath, str):
@@ -191,24 +202,24 @@ def get_inputs():
         module,
         init_args_fn_name,
         default=tuple(),
-        error_msg=f"Init args function '{init_args_fn_name}' not found in {filepath}"
+        error_msg=f"Init args function '{init_args_fn_name}' not found in {filepath}",
     )
     model_init_kwargs = maybe_load_and_run_callable(
         module,
         init_kwargs_fn_name,
         default={},
-        error_msg=f"Init kwargs function '{init_kwargs_fn_name}' not found in {filepath}"
+        error_msg=f"Init kwargs function '{init_kwargs_fn_name}' not found in {filepath}",
     )
     sample_args = load_and_run_callable(
         module,
         sample_args_fn_name,
-        f"Sample args function '{sample_args_fn_name}' not found in {filepath}"
+        f"Sample args function '{sample_args_fn_name}' not found in {filepath}",
     )
     sample_kwargs = maybe_load_and_run_callable(
         module,
         sample_kwargs_fn_name,
         default={},
-        error_msg=f"Sample kwargs function '{sample_kwargs_fn_name}' not found in {filepath}"
+        error_msg=f"Sample kwargs function '{sample_kwargs_fn_name}' not found in {filepath}",
     )
 
     nn_model: nn.Module = model(*model_init_args, **model_init_kwargs)
diff --git a/python/lighthouse/ingress/torch/utils.py b/python/lighthouse/ingress/torch/utils.py
@@ -43,8 +43,4 @@ def maybe_load_and_run_callable(
     """
     if symbol_name is None:
         return default
-    return load_and_run_callable(
-        module,
-        symbol_name,
-        error_msg=error_msg
-    )
+    return load_and_run_callable(module, symbol_name, error_msg=error_msg)
diff --git a/python/lighthouse/utils/__init__.py b/python/lighthouse/utils/__init__.py
diff --git a/python/lighthouse/utils/runtime_args.py b/python/lighthouse/utils/runtime_args.py

Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,5 @@`
`1`	`1`	`import sys`
	`2`	`+`
`2`	`3`	`from . import main`
`3`	`4`
`4`	`5`	# Invoke on command line with `python -m mlir_gen`.