Align functorch's flake8 config with pytorch's (#963)

Author: zou3519

.flake8

@@ -0,0 +1,14 @@
+[flake8]
+select = B,C,E,F,P,T4,W,B9
+max-line-length = 120
+# C408 ignored because we like the dict keyword argument syntax
+# E501 is not flexible enough, we're using B950 instead
+ignore =
+    E203,E305,E402,E501,E721,E741,F405,F821,F841,F999,W503,W504,C408,E302,W291,E303,
+    # shebang has extra meaning in fbcode lints, so I think it's not worth trying
+    # to line this up with executable bit
+    EXE001,
+    # these ignores are from flake8-bugbear; please fix!
+    B007,B008,
+    # these ignores are from flake8-comprehensions; please fix!
+    C400,C401,C402,C403,C404,C405,C407,C411,C413,C414,C415

codegen/gen.py

@@ -1,8 +1,5 @@
 import os
 import argparse
-import pathlib
-
-import torchgen
 from torchgen.gen import FileManager, parse_native_yaml
 from torchgen.gen import get_torchgen_root
 from gen_vmap_plumbing import gen_all_vmap_plumbing

examples/compilation/fuse_module.py

@@ -2,7 +2,6 @@
 from functorch.compile import compiled_module, tvm_compile
 import torch.nn as nn
 import torch
-from functools import partial
 
 
 def nop(f, _):

functorch/_src/partitioners.py

@@ -269,7 +269,7 @@ def classify_nodes(joint_module):
 
     aten = torch.ops.aten
 
-    pointwise_ops = [aten.add, aten.sub, aten.div, aten.atan2, aten.mul, aten.max, aten.min, aten.pow, aten.remainder, aten.fmod, aten.__and__, aten.__or__, aten.__xor__, aten.__lshift__, aten.__rshift__, aten.eq, aten.ne, aten.ge, aten.gt, aten.le, aten.lt, aten.abs, aten.bitwise_not, aten.ceil, aten.floor, aten.frac, aten.neg, aten.relu, aten.round, aten.silu, aten.trunc, aten.log, aten.log10, aten.log1p, aten.log2, aten.lgamma, aten.exp, aten.expm1, aten.erf, aten.erfc, aten.cos, aten.acos, aten.cosh, aten.sin, aten.asin, aten.sinh, aten.tan, aten.atan, aten.tanh, aten.atanh, aten.sqrt, aten.rsqrt,  aten.reciprocal, aten.sigmoid, aten.softplus, aten.threshold, aten.threshold_backward, aten.clamp, aten.where, aten.lerp, aten.addcmul, aten.gelu, aten.gelu_backward]  # noqa: E501
+    pointwise_ops = [aten.add, aten.sub, aten.div, aten.atan2, aten.mul, aten.max, aten.min, aten.pow, aten.remainder, aten.fmod, aten.__and__, aten.__or__, aten.__xor__, aten.__lshift__, aten.__rshift__, aten.eq, aten.ne, aten.ge, aten.gt, aten.le, aten.lt, aten.abs, aten.bitwise_not, aten.ceil, aten.floor, aten.frac, aten.neg, aten.relu, aten.round, aten.silu, aten.trunc, aten.log, aten.log10, aten.log1p, aten.log2, aten.lgamma, aten.exp, aten.expm1, aten.erf, aten.erfc, aten.cos, aten.acos, aten.cosh, aten.sin, aten.asin, aten.sinh, aten.tan, aten.atan, aten.tanh, aten.atanh, aten.sqrt, aten.rsqrt, aten.reciprocal, aten.sigmoid, aten.softplus, aten.threshold, aten.threshold_backward, aten.clamp, aten.where, aten.lerp, aten.addcmul, aten.gelu, aten.gelu_backward]  # noqa: E501
     misc_ops = [aten.to, aten.type_as, operator.getitem]
 
     reduction_ops = [aten.softmax, aten._softmax, aten._softmax_backward_data, aten.sum, aten.mean, aten._grad_sum_to_size, aten.sum_to_size, aten.amax]  # noqa: E501
@@ -338,27 +338,27 @@ def get_node_weight(node):
             continue
 
         if node in required_bw_nodes:
-            nx_graph.add_edge(node.name+"_in", "sink", capacity=math.inf)
+            nx_graph.add_edge(node.name + "_in", "sink", capacity=math.inf)
             continue
 
         if node.op == 'placeholder' and "primals" in node.target:
-            nx_graph.add_edge("source", node.name+"_in", capacity=math.inf)
+            nx_graph.add_edge("source", node.name + "_in", capacity=math.inf)
 
         # If a node can't be recomputed (too expensive or involves randomness),
         # we prevent it from being recomputed by adding an inf edge to the source
         # We only need to ban nodes in the fw pass, as those are the only ones that would be recomputed.
         if ban_recomputation(node) and node in required_fw_nodes:
-            nx_graph.add_edge("source", node.name+"_in", capacity=math.inf)
+            nx_graph.add_edge("source", node.name + "_in", capacity=math.inf)
 
         if 'tensor_meta' not in node.meta:
             weight = math.inf
         else:
             weight = get_node_weight(node)
 
         # Creates the weights on the "node" edge
-        nx_graph.add_edge(node.name+"_in", node.name+"_out", capacity=weight)
+        nx_graph.add_edge(node.name + "_in", node.name + "_out", capacity=weight)
         for user in node.users:
-            nx_graph.add_edge(node.name+"_out", user.name+"_in", capacity=math.inf)
+            nx_graph.add_edge(node.name + "_out", user.name + "_in", capacity=math.inf)
 
     cut_value, partition = nx.minimum_cut(nx_graph, "source", "sink")
     reachable, non_reachable = partition

test/common_utils.py

@@ -45,6 +45,7 @@ def loop(op, in_dims, out_dim, batch_size, *batched_args, **kwarg_values):
 # - Don't hash tensors in a global context, that'll keep them around forever
 def memoize(fn):
     memo = {}
+
     def wrapped(*args):
         if args not in memo:
             memo[args] = fn(*args)
@@ -78,7 +79,7 @@ def get_bdim_choices_batch_norm(num_tensors, _, running_mean=None, running_var=N
     options = (-1, None)
 
     # instance norm turns these into unbatched 0 tensors, so we cannot batch the input if either is not specified
-    if running_mean == None or running_var == None:
+    if running_mean is None or running_var is None:
         choices.append((None,) + (0,) * (num_tensors - 1))
         for choice in itertools.product(options, repeat=num_tensors - 1):
             choices.append((None,) + choice)

test/test_eager_transforms.py

@@ -365,7 +365,7 @@ def f(x):
         vjp_fn(result)
 
     def test_conj_bit(self):
-        x = torch.tensor(1+1j)
+        x = torch.tensor(1 + 1j)
 
         def foo(x):
             assert not x.is_conj()
@@ -1647,8 +1647,8 @@ def test_jacfwd_different_levels(self, device):
         A = 0.1 * torch.randn(b, d, d, device=device)
 
         def loss(A, x1, x2):
-            x2_hat = (A@(x1.T)).T
-            res = x2-x2_hat
+            x2_hat = (A @ (x1.T)).T
+            res = x2 - x2_hat
             res_sqr = res**2
             return res_sqr.sum()
 
@@ -2528,7 +2528,7 @@ def inner_loss(params, x1, y1):
             else:
                 loss = inner_loss(params, x1, y1)
                 grads = torch.autograd.grad(loss, params, create_graph=True)
-            new_params = [(params[i] - alpha*grads[i]) for i in range(len(params))]
+            new_params = [(params[i] - alpha * grads[i]) for i in range(len(params))]
 
             v_f = net(new_params, x2)
             return mse_loss(v_f, y2)
@@ -2537,15 +2537,15 @@ def inner_loss(params, x1, y1):
 
         # Compute with vmap+grad
         inner_losses = vmap(partial(get_loss_for_task, True))(task[0], task[1], task[2], task[3])
-        loss2 = sum(inner_losses)/len(inner_losses)
+        loss2 = sum(inner_losses) / len(inner_losses)
         result_grads = torch.autograd.grad(loss2, params)
 
         # Compute without vmap+grad
         inner_losses = [
             get_loss_for_task(False, task[0][i], task[1][i], task[2][i], task[3][i])
             for i in range(num_tasks)
         ]
-        loss2 = sum(inner_losses)/len(inner_losses)
+        loss2 = sum(inner_losses) / len(inner_losses)
         expected_grads = torch.autograd.grad(loss2, params)
 
         self.assertEqual(result_grads, expected_grads)

test/test_functionalize.py

@@ -1,12 +1,7 @@
-import torch
-
 import functorch
-from torch.testing._internal.common_utils import run_tests, TestCase, IS_WINDOWS
-import unittest
 from unittest.mock import patch
 import functools
-
-from functorch.compile import aot_function, nop
+from torch.testing._internal.common_utils import run_tests
 import test_compile_cache
 import test_pythonkey
 
@@ -41,7 +36,7 @@ class FunctionalizeTest(cls):
     return FunctionalizeTest
 
 
-FunctionalizeTestCompileCache = make_functionalize_test(test_compile_cache.TestCompileCache) 
+FunctionalizeTestCompileCache = make_functionalize_test(test_compile_cache.TestCompileCache)
 FunctionalizeTestCompileCacheStaticArgs = make_functionalize_test(test_compile_cache.TestCompileCacheStaticArgs)
 FunctionalizeTestPythonKeyAOT = make_functionalize_test(test_pythonkey.TestAOTAutograd)
 FunctionalizeTestPythonKeyContiguous = make_functionalize_test(test_pythonkey.TestContiguous)

test/test_memory_efficient_fusion.py

@@ -141,8 +141,8 @@ def layer_norm(x, weight, bias):
             mean = torch.mean(x, dim, keepdim=True)
             centered = x - mean
             var = torch.sum(centered * centered, dim, keepdim=True) / x.size(-1)
-            rvar = 1./torch.sqrt(var+eps)
-            normed = (x-mean) * rvar
+            rvar = 1. / torch.sqrt(var + eps)
+            normed = (x - mean) * rvar
             return normed * weight + bias
 
         bs = 10

test/test_minifier.py

@@ -2,7 +2,6 @@
 from functorch.compile import minifier
 from functorch import make_fx
 from torch.testing._internal.common_utils import TestCase, run_tests
-import unittest
 
 
 class TestMinifier(TestCase):

test/test_ops.py

@@ -13,7 +13,7 @@
 from torch.testing._internal.common_device_type import instantiate_device_type_tests
 from torch.testing._internal.common_device_type import ops
 from torch.testing._internal.common_device_type import \
-     toleranceOverride, tol
+    toleranceOverride, tol
 from functorch_lagging_op_db import functorch_lagging_op_db
 from functorch_additional_op_db import additional_op_db
 from common_utils import (
@@ -28,7 +28,6 @@
     opsToleranceOverride,
     check_vmap_fallback,
 )
-import unittest
 from torch.utils._pytree import tree_flatten, tree_unflatten, tree_map
 from functorch import grad, vjp, vmap, jacrev, jacfwd
 import torch.autograd.forward_ad as fwAD
@@ -773,8 +772,8 @@ def test_vmapjvpall(self, device, dtype, op):
         xfail('nn.functional.bilinear'),  # trilinear doesn't have batching rule
         xfail('linalg.eigh'),  # _linalg_eigh doesn't have batching rule
         xfail('linalg.eigvalsh'),  # _linalg_eigh doesn't have batching rule
-        xfail('logdet'), # _linalg_slogdet doesn't have batching rule
-        xfail('linalg.slogdet'), # _linalg_slogdet doesn't have batching rule
+        xfail('logdet'),  # _linalg_slogdet doesn't have batching rule
+        xfail('linalg.slogdet'),  # _linalg_slogdet doesn't have batching rule
     }))
     @toleranceOverride({torch.float32: tol(atol=1e-04, rtol=1e-04)})
     def test_vmapjvpall_has_batch_rule(self, device, dtype, op):
@@ -1133,6 +1132,7 @@ def reference(primals, cotangents, primals_tangents, cotangents_tangents):
                 self.assertFalse(op.supports_fwgrad_bwgrad,
                                  f"{op.name} now supports forward over reverse without a decomposition. " +
                                  "Please remove the decomposition version")
+
                 def is_differentiable(t):
                     return isinstance(t, torch.Tensor) and t.dtype == torch.float32
                 args = (cotangents, *primals)
@@ -1148,7 +1148,7 @@ def is_differentiable(t):
                 self.assertEqual(result, expected)
 
     def _make_extremal_inputs(self, shape, device):
-        if shape == None:
+        if shape is None:
             return (None,)
         return (
             torch.full(shape, -1000., device=device),