Use ruff to format and lint Python code (#8684)

Author: alexreinking

.github/workflows/presubmit.yml

@@ -17,21 +17,25 @@ permissions:
 
 jobs:
   check_clang_format:
-    name: Check clang-format
+    name: Check clang-format and ruff
     runs-on: macos-14
     steps:
-      - uses: actions/checkout@v3
+      - uses: actions/checkout@v4
+
       - name: Install clang-format
         run: brew install llvm@19
+
       - name: Check clang-format
         run: ./run-clang-format.sh -c
         env:
           CLANG_FORMAT_LLVM_INSTALL_DIR: /opt/homebrew/opt/llvm@19
+
+      - uses: astral-sh/ruff-action@v3
   check_clang_tidy:
     name: Check clang-tidy
     runs-on: macos-14
     steps:
-      - uses: actions/checkout@v3
+      - uses: actions/checkout@v4
       - name: Install clang-tidy
         run: brew install llvm@19 ninja lld@19
       - name: Run clang-tidy
@@ -42,7 +46,7 @@ jobs:
     name: Check CMake file lists
     runs-on: ubuntu-22.04
     steps:
-      - uses: actions/checkout@v3
+      - uses: actions/checkout@v4
       - name: Run test sources check
         run: |
           shopt -s nullglob

apps/HelloPyTorch/modules.py

@@ -7,6 +7,7 @@
 # TODO(mgharbi): maybe find a way to wrap function and module directly in C++
 # instead of generating the C++ wrapper on the fly?
 
+
 def _dispatch(opname, optype=th.float32, cuda=False):
     """
     Helper function that matches an opname and type to the Halide backend.
@@ -23,7 +24,7 @@ def _dispatch(opname, optype=th.float32, cuda=False):
       op: a python function wrapping the requested Halide operator.
     """
 
-    assert type(opname) == str, "opname should be a string"
+    assert type(opname) is str, "opname should be a string"
     assert type(optype) == th.dtype, "optype should be a tensor datatype (torch.dtype)"
 
     if cuda:
@@ -34,12 +35,13 @@ def _dispatch(opname, optype=th.float32, cuda=False):
     elif optype == th.float64:
         opname += "_float64"
     else:
-        raise ValueError("Optype %s not recognized %s" % optype)
+        raise ValueError("Optype {} not recognized {}".format(*optype))
     op = getattr(ops, opname)
     if not hasattr(ops, opname):
-        raise ValueError("Module has no operator %s" % opname)
+        raise ValueError(f"Module has no operator {opname}")
     return op
 
+
 def _forward_common(ctx, input_a, input_b):
     tp = input_a.dtype
     cuda = input_a.is_cuda
@@ -57,6 +59,7 @@ def _forward_common(ctx, input_a, input_b):
     fn_(input_a, input_b, out)
     return out
 
+
 def _backward_common(ctx, d_out, backward_op):
     tp = d_out.dtype
     cuda = d_out.is_cuda
@@ -75,48 +78,54 @@ def _backward_common(ctx, d_out, backward_op):
     fn_(input_a, input_b, d_out.contiguous(), d_input_a, d_input_b)
     return d_input_a, d_input_b
 
+
 # TODO(srj): surely there's a better way to do this,
 # but PyTorch seems to make it tricky to pass in
 # extra info to the backward() method.
 class AddFunction_Grad(th.autograd.Function):
-  """Version using the manually-written backprop"""
-  def __init__(self):
-      super(AddFunction_Grad, self).__init__()
+    """Version using the manually-written backprop"""
+
+    def __init__(self):
+        super().__init__()
 
-  @staticmethod
-  def forward(ctx, input_a, input_b):
-      return _forward_common(ctx, input_a, input_b)
+    @staticmethod
+    def forward(ctx, input_a, input_b):
+        return _forward_common(ctx, input_a, input_b)
+
+    @staticmethod
+    def backward(ctx, d_out):
+        return _backward_common(ctx, d_out, "add_grad")
 
-  @staticmethod
-  def backward(ctx, d_out):
-      return _backward_common(ctx, d_out, "add_grad")
 
 class AddFunction_HalideGrad(th.autograd.Function):
-  """Version using the Halide-generated backprop"""
-  def __init__(self):
-      super(AddFunction_HalideGrad, self).__init__()
+    """Version using the Halide-generated backprop"""
+
+    def __init__(self):
+        super().__init__()
 
-  @staticmethod
-  def forward(ctx, input_a, input_b):
-      return _forward_common(ctx, input_a, input_b)
+    @staticmethod
+    def forward(ctx, input_a, input_b):
+        return _forward_common(ctx, input_a, input_b)
+
+    @staticmethod
+    def backward(ctx, d_out):
+        return _backward_common(ctx, d_out, "add_halidegrad")
 
-  @staticmethod
-  def backward(ctx, d_out):
-      return _backward_common(ctx, d_out, "add_halidegrad")
 
 class Add(th.nn.Module):
     """Defines a module that uses our autograd function.
 
     This is so we can use it as an operator.
     """
+
     def __init__(self, backward_op):
-        super(Add, self).__init__()
+        super().__init__()
         if backward_op == "add_grad":
-          self._adder = AddFunction_Grad
+            self._adder = AddFunction_Grad
         elif backward_op == "add_halidegrad":
-          self._adder = AddFunction_HalideGrad
+            self._adder = AddFunction_HalideGrad
         else:
-          assert False
+            assert False
 
     def forward(self, a, b):
         return self._adder.apply(a, b)

apps/HelloPyTorch/setup.py

@@ -1,41 +1,41 @@
 """Synthesizes the cpp wrapper code and builds dynamic Python extension."""
+
 import os
 import platform
 import re
-from setuptools import setup, find_packages
+from setuptools import setup
 
 from torch.utils.cpp_extension import BuildExtension
-import torch as th
 
 
 def generate_pybind_wrapper(path, headers, has_cuda):
-    s = "#include \"torch/extension.h\"\n\n"
+    s = '#include "torch/extension.h"\n\n'
     if has_cuda:
-        s += "#include \"HalidePyTorchCudaHelpers.h\"\n"
-    s += "#include \"HalidePyTorchHelpers.h\"\n"
+        s += '#include "HalidePyTorchCudaHelpers.h"\n'
+    s += '#include "HalidePyTorchHelpers.h"\n'
     for h in headers:
-        s += "#include \"{}\"\n".format(os.path.splitext(h)[0]+".pytorch.h")
+        s += '#include "{}"\n'.format(os.path.splitext(h)[0] + ".pytorch.h")
 
     s += "\nPYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {\n"
     for h in headers:
         name = os.path.splitext(h)[0]
-        s += "  m.def(\"{}\", &{}_th_, \"PyTorch wrapper of the Halide pipeline {}\");\n".format(
-          name, name, name)
+        s += f'  m.def("{name}", &{name}_th_, "PyTorch wrapper of the Halide pipeline {name}");\n'
     s += "}\n"
-    with open(path, 'w') as fid:
+    with open(path, "w") as fid:
         fid.write(s)
 
+
 if __name__ == "__main__":
     # This is where the generate Halide ops headers live. We also generate the .cpp
     # wrapper in this directory
     build_dir = os.getenv("BIN")
     if build_dir is None or not os.path.exists(build_dir):
-        raise ValueError("Bin directory {} is invalid".format(build_dir))
+        raise ValueError(f"Bin directory {build_dir} is invalid")
 
     # Path to a distribution of Halide
     halide_dir = os.getenv("HALIDE_DISTRIB_PATH")
     if halide_dir is None or not os.path.exists(halide_dir):
-        raise ValueError("Halide directory {} is invalid".format(halide_dir))
+        raise ValueError(f"Halide directory {halide_dir} is invalid")
 
     has_cuda = os.getenv("HAS_CUDA")
     if has_cuda is None or has_cuda == "0":
@@ -76,27 +76,38 @@ def generate_pybind_wrapper(path, headers, has_cuda):
         print("Generating CUDA wrapper")
         generate_pybind_wrapper(wrapper_path, hl_headers, True)
         from torch.utils.cpp_extension import CUDAExtension
-        extension = CUDAExtension(ext_name, sources,
-                                  include_dirs=include_dirs,
-                                  extra_objects=hl_libs,
-                                  libraries=["cuda"],  # Halide ops need the full cuda lib, not just the RT library
-                                  extra_compile_args=compile_args)
+
+        extension = CUDAExtension(
+            ext_name,
+            sources,
+            include_dirs=include_dirs,
+            extra_objects=hl_libs,
+            libraries=[
+                "cuda"
+            ],  # Halide ops need the full cuda lib, not just the RT library
+            extra_compile_args=compile_args,
+        )
     else:
         print("Generating CPU wrapper")
         generate_pybind_wrapper(wrapper_path, hl_headers, False)
         from torch.utils.cpp_extension import CppExtension
-        extension = CppExtension(ext_name, sources,
-                                 include_dirs=include_dirs,
-                                 extra_objects=hl_libs,
-                                 extra_compile_args=compile_args)
+
+        extension = CppExtension(
+            ext_name,
+            sources,
+            include_dirs=include_dirs,
+            extra_objects=hl_libs,
+            extra_compile_args=compile_args,
+        )
 
     # Build the Python extension module
-    setup(name=ext_name,
-          verbose=True,
-          url="",
-          author_email="[email protected]",
-          author="Some Author",
-          version="0.0.0",
-          ext_modules=[extension],
-          cmdclass={"build_ext": BuildExtension}
-          )
+    setup(
+        name=ext_name,
+        verbose=True,
+        url="",
+        author_email="[email protected]",
+        author="Some Author",
+        version="0.0.0",
+        ext_modules=[extension],
+        cmdclass={"build_ext": BuildExtension},
+    )

apps/HelloPyTorch/test.py

@@ -1,7 +1,5 @@
 """Verifies the Halide operator functions properly."""
 
-
-import os
 import unittest
 import warnings
 
@@ -12,8 +10,8 @@
 class TestAdd(unittest.TestCase):
     def setUp(self):
         self.a = th.ones(1, 2, 8, 8)
-        self.b = th.ones(1, 2, 8, 8)*3
-        self.gt = th.ones(1, 2, 8, 8)*4
+        self.b = th.ones(1, 2, 8, 8) * 3
+        self.gt = th.ones(1, 2, 8, 8) * 4
 
     def test_cpu_single(self):
         self._test_add(is_double=False)
@@ -53,23 +51,22 @@ def _test_add(self, is_cuda=False, is_double=False):
             else:
                 print("  .Single-precision mode, backward_op:", backward_op)
 
-            diff = (output-self.gt).sum().item()
-            assert diff == 0.0, "Test failed: sum should be 4, got %f" % diff
+            diff = (output - self.gt).sum().item()
+            assert diff == 0.0, f"Test failed: sum should be 4, got {diff:f}"
 
             self.a.requires_grad = True
             self.b.requires_grad = True
 
             for i in range(100):
                 output = add(self.a, self.b).sum()
                 output.backward()
-            
+
             # Inputs are float, the gradient checker wants double inputs and
             # will issue a warning.
-            warnings.filterwarnings(
-                "ignore", module=r".*gradcheck*")
+            warnings.filterwarnings("ignore", module=r".*gradcheck*")
 
             # Test the gradient is correct
-            res = th.autograd.gradcheck(add, [self.a, self.b], eps=1e-2)
+            th.autograd.gradcheck(add, [self.a, self.b], eps=1e-2)
 
             print("     Test ran successfully: difference is", diff)
 

apps/depthwise_separable_conv/tf_separable_conv.py

@@ -1,22 +1,26 @@
 import tensorflow as tf
 import time
 
-with tf.device('/GPU:0'):
-
+with tf.device("/GPU:0"):
     img = tf.random.uniform([4, 112, 112, 32])
     depthwise_filter = tf.random.uniform([3, 3, 32, 1])
     pointwise_filter = tf.random.uniform([1, 1, 32 * 1, 16])
-    
+
     best = None
     num_trials = 10
     num_iter = 10
     for j in range(num_trials):
         start = time.time()
         for i in range(num_iter):
             out = tf.nn.separable_conv2d(
-                img, depthwise_filter, pointwise_filter,
-                strides = (1, 1, 1, 1), padding = 'VALID')
+                img,
+                depthwise_filter,
+                pointwise_filter,
+                strides=(1, 1, 1, 1),
+                padding="VALID",
+            )
         end = time.time()
         t = (end - start) / num_iter
-        if not best or t < best: best = t
-    print('time: {} ms'.format(1000 * best))
+        if not best or t < best:
+            best = t
+    print(f"time: {1000 * best} ms")