format

Author: makslevental

tests/test_arith.py

@@ -1,11 +1,8 @@
-from textwrap import dedent
-
 import mlir.extras.types as T
 import pytest
 
 from mlir.extras.ast.canonicalize import canonicalize
 from mlir.extras.dialects.ext import arith
-from mlir.extras.dialects.ext.arith import Scalar
 from mlir.extras.dialects.ext.func import func
 
 # noinspection PyUnresolvedReferences

tests/test_async.py

@@ -1,8 +1,7 @@
 import platform
-
-import numpy as np
 from textwrap import dedent
 
+import numpy as np
 import pytest
 
 from mlir.extras.runtime.passes import Pipeline

tests/test_func.py

@@ -1,19 +1,17 @@
 import inspect
 import sys
 import threading
-from textwrap import dedent
 from typing import TypeVar
 
-import pytest
-
 import mlir.extras.types as T
+import pytest
+from mlir.ir import FunctionType
 
 from mlir.extras.ast.canonicalize import canonicalize
 from mlir.extras.context import mlir_mod_ctx, RAIIMLIRContextModule
+from mlir.extras.dialects.ext import linalg, arith, scf
 from mlir.extras.dialects.ext.arith import constant
 from mlir.extras.dialects.ext.func import func
-from mlir.extras.dialects.ext import linalg, arith, scf, memref
-from mlir.ir import FunctionType
 
 # noinspection PyUnresolvedReferences
 from mlir.extras.testing import (
@@ -171,7 +169,6 @@ def test_func_no_context_2(ctx: MLIRContext):
 
 
 def test_generics_just_args(ctx: MLIRContext):
-
     @func(generics=generics)
     def mat_product_kernel(
         A: "T.memref(M, K, dtype)",
@@ -212,7 +209,6 @@ def mat_product_kernel(
 
 
 def test_generics_with_canonicalizations(ctx: MLIRContext):
-
     generics = M, K, N, dtype = list(map(TypeVar, ["M", "K", "N", "dtype"]))
 
     @func(generics=generics)

tests/test_gpu.py

@@ -710,7 +710,6 @@ def _():
 
 
 def test_amdgpu(ctx: MLIRContext):
-
     set_container_module(ctx.module)
 
     M, K, N, dtype = 32, 32, 32, T.f32()

tests/test_linalg.py

@@ -1,8 +1,6 @@
-from textwrap import dedent
-
+import mlir.extras.types as T
 import pytest
 
-import mlir.extras.types as T
 from mlir.extras.dialects.ext import linalg, memref, tensor
 
 # noinspection PyUnresolvedReferences

tests/test_llvm.py

@@ -1,8 +1,8 @@
 from textwrap import dedent
 
+import mlir.extras.types as T
 import pytest
 
-import mlir.extras.types as T
 from mlir.extras.dialects.ext import llvm
 from mlir.extras.dialects.ext.func import func
 

tests/test_nvgpu_nvvm.py

@@ -418,7 +418,6 @@ def main(module: any_op_t()):
 
 CUDA_RUNTIME_LIB_PATH = Path(_mlir_libs.__file__).parent / f"libmlir_cuda_runtime.so"
 
-
 NVIDIA_GPU = False
 try:
     subprocess.check_output("nvidia-smi")
@@ -669,7 +668,6 @@ def main(module: any_op_t()):
 
 
 def test_tma(ctx: MLIRContext):
-
     M = K = N = 64
 
     @gpu.func

tests/test_other_hosts.py

@@ -1,11 +1,10 @@
 import pytest
 
-from util import jax_not_installed, mlir_bindings_not_installed
+from util import jax_not_installed
 
 
 @pytest.mark.skipif(jax_not_installed(), reason="jax not installed")
 def test_jax_trampolines_smoke():
-    from mlir import ir
     from jaxlib.mlir import ir
 
     # noinspection PyUnresolvedReferences

tests/test_regions.py

@@ -1,21 +1,19 @@
-from textwrap import dedent
-
-import pytest
-
 import mlir.extras.types as T
+import pytest
+from mlir.dialects.builtin import module
 from mlir.dialects.func import return_
 from mlir.dialects.memref import alloca_scope, alloca_scope_return
 from mlir.dialects.scf import yield_ as scf_yield
 from mlir.dialects.tensor import rank, yield_ as tensor_yield
-from mlir.dialects.builtin import module
+from mlir.extras.types import tensor
+
 from mlir.extras.dialects.ext import linalg, memref
 from mlir.extras.dialects.ext.arith import constant
 from mlir.extras.dialects.ext.cf import br, cond_br
 from mlir.extras.dialects.ext.func import func
 from mlir.extras.dialects.ext.memref import alloca_scope
 from mlir.extras.dialects.ext.scf import execute_region
 from mlir.extras.dialects.ext.tensor import S, generate
-from mlir.extras.util import bb
 
 # noinspection PyUnresolvedReferences
 from mlir.extras.testing import (
@@ -24,7 +22,7 @@
     filecheck_with_comments,
     MLIRContext,
 )
-from mlir.extras.types import tensor
+from mlir.extras.util import bb
 
 # needed since the fix isn't defined here nor conftest.py
 pytest.mark.usefixtures("ctx")
@@ -409,7 +407,6 @@ def matmul_i16_i16(
 
 
 def test_defer_emit_1(ctx: MLIRContext):
-
     matmul_i16_i16.emit(decl=True)
 
     @module
@@ -428,7 +425,6 @@ def mod():
 
 
 def test_defer_emit_2(ctx: MLIRContext):
-
     matmul_i16_i16.emit(force=True)
 
     @module
@@ -459,7 +455,6 @@ def matmul_i16_i16(
 
 
 def test_defer_emit_3(ctx: MLIRContext):
-
     matmul_i16_i16.emit(force=True)
 
     @module

tests/test_runtime.py

@@ -3,24 +3,24 @@
 import re
 from textwrap import dedent
 
+import mlir.extras.types as T
 import numpy as np
 import pytest
+from mlir.dialects.arith import sitofp, index_cast
+from mlir.dialects.memref import cast
 from mlir.ir import UnitAttr, Module, StridedLayoutAttr, InsertionPoint, Context
 from mlir.runtime import get_unranked_memref_descriptor, get_ranked_memref_descriptor
 
-import mlir.extras.types as T
 from mlir.extras.ast.canonicalize import canonicalize
 from mlir.extras.context import RAIIMLIRContext, ExplicitlyManagedModule
 from mlir.extras.dialects.ext import linalg
-from mlir.dialects.arith import sitofp, index_cast
 from mlir.extras.dialects.ext.arith import constant
 from mlir.extras.dialects.ext.func import func
 from mlir.extras.dialects.ext.memref import load, store, S
 from mlir.extras.dialects.ext.scf import (
     canonicalizer,
     range_,
 )
-from mlir.dialects.memref import cast
 from mlir.extras.runtime.passes import Pipeline, run_pipeline
 from mlir.extras.runtime.refbackend import (
     LLVMJITBackend,