Repository: pytorch/functorch (language: python)

Issue:
Set up some linters
Like flake8


Discussion / hints:
@zou3519 how about using black ?

Fix patch:
diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml
new file mode 100644
index 000000000..3a0f63b77
--- /dev/null
+++ b/.github/workflows/lint.yml
@@ -0,0 +1,31 @@
+name: Lint
+
+on:
+  push:
+    branches:
+      - main
+  pull_request:
+
+jobs:
+  flake8-py3:
+    runs-on: ubuntu-18.04
+    steps:
+      - uses: actions/checkout@v2
+      - name: Setup Python
+        uses: actions/setup-python@v2
+        with:
+          python-version: 3.x
+          architecture: x64
+      - name: Install dependencies
+        run: |
+          set -eux
+          pip3 install flake8 --user
+          flake8 --version
+      - name: Run flake8
+        run: |
+          set -eux
+          flake8 .
+
+concurrency:
+  group: lint-${{ github.event.pull_request.number || github.sha }}-${{ github.event_name == 'workflow_dispatch' }}
+  cancel-in-progress: true
diff --git a/codegen/codegen_outofplacebatching.py b/codegen/codegen_outofplacebatching.py
index 63a018c19..8e39fd7a8 100644
--- a/codegen/codegen_outofplacebatching.py
+++ b/codegen/codegen_outofplacebatching.py
@@ -6,30 +6,34 @@
 
 import argparse
 from typing import Tuple, List
-from collections import defaultdict
 import re
 
+
 def num_leading_spaces(line: str) -> int:
     return len(line) - len(line.lstrip())
 
+
 def min_leading_spaces(lines):
     num_spaces = [num_leading_spaces(line) for line in lines if len(line) > 0]
     if len(num_spaces) == 0:
         return None
     return min(num_spaces)
 
+
 def deindent(code: str) -> str:
     lines = code.split('\n')
     mls = min_leading_spaces(lines)
     lines = [line[mls:] for line in lines]
     return '\n'.join(lines)
 
+
 def indent(code: str, num) -> str:
     lines = code.split('\n')
     indented_lines = [' ' * num + line for line in lines]
     indented_lines[0] = lines[0]
     return '\n'.join(indented_lines)
 
+
 def is_tensor(typ: str) -> bool:
     if typ == 'Tensor':
         return True
@@ -37,6 +41,7 @@ def is_tensor(typ: str) -> bool:
         return True
     return False
 
+
 def is_optional_tensor(typ: str) -> bool:
     if typ == 'c10::optional<Tensor>':
         return True
@@ -44,10 +49,12 @@ def is_optional_tensor(typ: str) -> bool:
         return True
     return False
 
+
 def is_vector_tensor(typ: str) -> bool:
     # (chilli): I don't really understand why there's 2 dots in front?
     return (typ == '::std::vector<Tensor>')
 
+
 def add_bdim_after_tensor(types: Tuple[str]) -> Tuple[str]:
     result = []
     for typ in types:
@@ -56,6 +63,7 @@ def add_bdim_after_tensor(types: Tuple[str]) -> Tuple[str]:
             result.append('c10::optional<int64_t>')
     return tuple(result)
 
+
 def batch_rule_type(
         op_returns: Tuple[str],
         op_args: Tuple[str],
@@ -67,6 +75,7 @@ def batch_rule_type(
     result = f"typedef std::tuple<{','.join(returns)}> (*{br_t})({', '.join(args)});"
     return result, br_t
 
+
 def unwrap_tensor(name: str) -> List[str]:
     result = f"""\
     Tensor {name}_value;
@@ -74,6 +83,7 @@ def unwrap_tensor(name: str) -> List[str]:
     std::tie({name}_value, {name}_bdim) = unwrapTensorAtLevel({name}, cur_level);"""
     return deindent(result).split('\n')
 
+
 def unwrap_optional_tensor(name: str) -> List[str]:
     result = f"""\
     optional<Tensor> {name}_value;
@@ -83,6 +93,7 @@ def unwrap_optional_tensor(name: str) -> List[str]:
     }}"""
     return deindent(result).split('\n')
 
+
 def gen_unwraps(arg_types, arg_names):
     tensors = [name for typ, name in zip(arg_types, arg_names) if is_tensor(typ)]
     optional_tensors = [name for typ, name in zip(arg_types, arg_names) if is_optional_tensor(typ)]
@@ -103,6 +114,7 @@ def gen_unwraps(arg_types, arg_names):
             unwrapped_arg_list.append(arg)
     return unwraps, unwrapped_arg_list
 
+
 def lower(returns: Tuple[str], args: List[Tuple[str, str]], unique_count: int, ops) -> str:
     arg_types, arg_names = zip(*args)
     batch_rule_typedef, batch_rule_t = batch_rule_type(returns, arg_types, unique_count)
@@ -120,7 +132,9 @@ def lower(returns: Tuple[str], args: List[Tuple[str, str]], unique_count: int, o
             wrapped_returns.append(f'makeBatched(std::get<{idx}>(results), std::get<{idx + 1}>(results), cur_level)')
             idx += 2
         elif is_vector_tensor(ret):
-            wrapped_returns.append(f'makeBatchedVector(std::get<{idx}>(results), std::get<{idx + 1}>(results), cur_level)')
+            wrapped_returns.append(
+                f'makeBatchedVector(std::get<{idx}>(results), std::get<{idx + 1}>(results), cur_level)'
+            )
             idx += 2
         else:
             wrapped_returns.append(f'std::get<{idx}>(results)')
@@ -148,6 +162,7 @@ def lower(returns: Tuple[str], args: List[Tuple[str, str]], unique_count: int, o
     }}"""
     return deindent(result)
 
+
 def parse_return(return_t):
     if 'std::tuple' not in return_t:
         return (return_t,)
@@ -156,6 +171,7 @@ def parse_return(return_t):
         m = re.match(r'::std::tuple<(.*)>', return_t)
     return tuple([x.strip() for x in m.group(1).split(',')])
 
+
 def parse_args(args_t):
     # There is an assumption made that args are separated with comma-space
     # and types like std::array<bool,2> do not contain spaces after the comma
@@ -166,6 +182,7 @@ def parse_args(args_t):
         result.append((arg[:split_idx].strip(), arg[split_idx:].strip()))
     return tuple(result)
 
+
 def get_signatures(path='build/aten/src/ATen/RegistrationDeclarations.h', include_op=False):
     with open(path, 'r') as f:
         txt = f.read()
@@ -188,6 +205,7 @@ def get_signatures(path='build/aten/src/ATen/RegistrationDeclarations.h', includ
         schemas.append(result)
     return tuple(schemas)
 
+
 def is_schema_outplace(schema):
     _, returns, args = schema
     for arg in args:
@@ -207,16 +225,19 @@ def is_schema_outplace(schema):
             return False
     return True
 
+
 def get_hash(schema):
     ret_t, args = schema
     args_t, _ = tuple(zip(*args))
     return (ret_t, args_t)
 
+
 class Container:
     def __init__(self, schema, ops):
         self.schema = schema
         self.ops = ops
 
+
 if __name__ == '__main__':
     parser = argparse.ArgumentParser()
     parser.add_argument('path',
diff --git a/codegen/gen_plumbing.py b/codegen/gen_plumbing.py
index 9b785a337..bbf1a510a 100644
--- a/codegen/gen_plumbing.py
+++ b/codegen/gen_plumbing.py
@@ -1,5 +1,4 @@
 import argparse
-import re
 
 from codegen_outofplacebatching import deindent, get_signatures, gen_unwraps
 
diff --git a/examples/compilation/eager_fusion.py b/examples/compilation/eager_fusion.py
index 7709a2d1f..721b73f20 100644
--- a/examples/compilation/eager_fusion.py
+++ b/examples/compilation/eager_fusion.py
@@ -1,5 +1,4 @@
 from functorch import compiled_function, tvm_compile
-import torch.nn as nn
 import torch
 from functools import partial
 import time
@@ -7,9 +6,12 @@
 
 a = torch.randn(2000, 1, 4, requires_grad=True)
 b = torch.randn(1, 2000, 4)
+
+
 def f(a):
     return (a * b).sum(dim=0)
 
+
 fw_compiler = partial(tvm_compile, name='fw_keops')
 bw_compiler = partial(tvm_compile, name='bw_keops')
 compiled_f = compiled_function(f, fw_compiler, bw_compiler)
@@ -19,30 +21,35 @@ def f(a):
 iters = 10
 out = compiled_f(a)
 out.sum().backward()
+
+
 def bench(func):
     begin = time.time()
     for _ in range(iters):
         out = func(a).sin()
         out.sum().backward()
         a.grad = None
-    print(time.time()-begin)
+    print(time.time() - begin)
+
 
 def bench_jax():
     import jax.numpy as jnp
     import jax
     jax_a = jnp.array(a.detach().numpy())
     jax_b = jnp.array(b.detach().numpy())
+
     def f(a):
-        return jnp.sin((a*jax_b).sum(axis=[0])).sum()
+        return jnp.sin((a * jax_b).sum(axis=[0])).sum()
     jit_f = jax.jit(jax.grad(f))
     jit_f(jax_a)
     begin = time.time()
     for _ in range(iters):
         out = jit_f(jax_a)
     out.block_until_ready()
-    print(time.time()-begin)
+    print(time.time() - begin)
     # for
 
+
 bench(f)
 bench(compiled_f)
 # bench_jax()
diff --git a/examples/compilation/fuse_module.py b/examples/compilation/fuse_module.py
index aacac37a8..627484b43 100644
--- a/examples/compilation/fuse_module.py
+++ b/examples/compilation/fuse_module.py
@@ -1,22 +1,27 @@
+import timeit
 from functorch import compiled_module, tvm_compile
 import torch.nn as nn
 import torch
 from functools import partial
 
+
 def nop(f, _):
     return f
 
+
 fw_compiler = partial(tvm_compile, name='fw_keops')
 bw_compiler = partial(tvm_compile, name='bw_keops')
 fw_compiler = nop
 bw_compiler = nop
 
+
 def run(mod, input):
     out = mod(input)
     out.sum().backward()
     grads = [p.grad for p in mod.parameters()]
     return (out, *grads)
 
+
 class Foo(nn.Module):
     def __init__(self):
         super(Foo, self).__init__()
@@ -26,6 +31,7 @@ def __init__(self):
     def forward(self, x):
         return (self.param * x + self.buf).sum(dim=0)
 
+
 input = torch.randn(1)
 mod = Foo()
 compiled_mod = compiled_module(mod, fw_compiler, bw_compiler)
@@ -42,7 +48,6 @@ def forward(self, x):
 for a, b in zip(run(mod, input), run(compiled_mod, input)):
     torch.testing.assert_allclose(a, b)
 
-import timeit
 for _ in range(5):
     i = 10000
     t = timeit.Timer("mod(input)", globals=globals()).timeit(10000)
diff --git a/examples/compilation/linear_train.py b/examples/compilation/linear_train.py
index 9ed01b7c4..ac9b9e60b 100644
--- a/examples/compilation/linear_train.py
+++ b/examples/compilation/linear_train.py
@@ -4,20 +4,20 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-from functorch import grad, vmap, pythonkey_trace, wrap_key, make_fx, nnc_jit, make_functional, grad_and_value
+from functorch import nnc_jit, make_functional
 import torch
-import torch.fx as fx
 import torch.nn as nn
 import time
 torch._C._jit_override_can_fuse_on_cpu(True)
 
+
 def bench(f, iters=100, warmup=10):
     for _ in range(warmup):
         f()
     begin = time.time()
     for _ in range(iters):
         f()
-    print((time.time()-begin))
+    print((time.time() - begin))
 
 
 class Foo(nn.Module):
@@ -42,16 +42,20 @@ def forward(self, x):
 jit_mod = torch.jit.script(mod)
 
 func_model, weights = make_functional(mod)
-lr =1.0
+lr = 1.0
+
 
 def functional_step(x, weights):
     weights = [weight.detach().requires_grad_() for weight in weights]
     out = func_model(weights, x)
     out.backward()
-    new_weights = [weight - lr*weight.grad for weight in weights]
+    new_weights = [weight - lr * weight.grad for weight in weights]
     return out, new_weights
 
-optim = torch.optim.SGD(jit_mod.parameters(), lr=lr, momentum=0,dampening=0, weight_decay=0)
+
+optim = torch.optim.SGD(jit_mod.parameters(), lr=lr, momentum=0, dampening=0, weight_decay=0)
+
+
 def jit_step(x, weights):
     optim.zero_grad()
     loss = jit_mod(x)
@@ -59,6 +63,7 @@ def jit_step(x, weights):
     optim.step()
     return loss, None
 
+
 def train(train_step, weights):
     torch.manual_seed(16)
     train_step(inp, weights)
@@ -67,9 +72,10 @@ def train(train_step, weights):
         loss, weights = train_step(torch.randn(batch_size, features), weights)
         if itr % 200 == 0:
             print(f"Loss at {itr}: {loss}")
-    print("Time taken: ", time.time()-begin)
+    print("Time taken: ", time.time() - begin)
     print()
 
+
 grad_pt = functional_step
 grad_nnc = nnc_jit(functional_step)
 
diff --git a/examples/compilation/simple_function.py b/examples/compilation/simple_function.py
index d0eb1b385..791990b1a 100644
--- a/examples/compilation/simple_function.py
+++ b/examples/compilation/simple_function.py
@@ -8,9 +8,11 @@
 import torch
 import time
 
+
 def f(x):
     return torch.sin(x).sum()
 
+
 inp = torch.randn(100)
 grad_pt = grad(f)
 grad_fx = make_fx(grad_pt)(inp)
@@ -18,6 +20,7 @@ def f(x):
 loopnest = make_nnc(grad_pt)(inp)
 print(loopnest)
 
+
 def bench(name, f, iters=10000, warmup=3):
     for _ in range(warmup):
         f()
@@ -26,6 +29,7 @@ def bench(name, f, iters=10000, warmup=3):
         f()
     print(f"{name}: ", time.time() - begin)
 
+
 bench("Pytorch: ", lambda: grad_pt(inp))
 bench("FX: ", lambda: grad_fx(inp))
 bench("NNC: ", lambda: grad_nnc(inp))
diff --git a/examples/dp_cifar10/cifar10_transforms.py b/examples/dp_cifar10/cifar10_transforms.py
index 1a7898014..11a9a6a85 100644
--- a/examples/dp_cifar10/cifar10_transforms.py
+++ b/examples/dp_cifar10/cifar10_transforms.py
@@ -24,7 +24,6 @@
 from torchvision.datasets import CIFAR10
 from tqdm import tqdm
 
-from functools import partial
 import functorch
 from functorch import vmap, grad_and_value
 from functorch import make_functional
@@ -32,6 +31,7 @@
 # disable warning spam
 functorch._C._set_vmap_fallback_warning_enabled(False)
 
+
 def save_checkpoint(state, is_best, filename="checkpoint.tar"):
     torch.save(state, filename)
     if is_best:
@@ -48,6 +48,7 @@ def compute_norms(sample_grads):
     norms = torch.stack(norms, dim=0).norm(2, dim=0)
     return norms
 
+
 def clip_and_accumulate_and_add_noise(model, max_per_sample_grad_norm=1.0, noise_multiplier=1.0):
     sample_grads = tuple(param.grad_sample for param in model.parameters())
 
@@ -73,6 +74,7 @@ def clip_and_accumulate_and_add_noise(model, max_per_sample_grad_norm=1.0, noise
         param.grad = param_grad
         del param.grad_sample
 
+
 def train(args, model, train_loader, optimizer, epoch, device):
     model.train()
     criterion = nn.CrossEntropyLoss()
@@ -106,10 +108,10 @@ def compute_loss_and_output(weights, image, target):
         # computes gradients by running both the forward and backward pass.
         # We want to extract some intermediate
         # values from the computation (i.e. the loss and output).
-        # 
+        #
         # To extract the loss, we use the `grad_and_value` API, that returns the
         # gradient of the weights w.r.t. the loss and the loss.
-        # 
+        #
         # To extract the output, we use the `has_aux=True` flag.
         # `has_aux=True` assumes that `f` returns a tuple of two values,
         # where the first is to be differentiated and the second "auxiliary value"
@@ -152,6 +154,7 @@ def compute_loss_and_output(weights, image, target):
                 f"Acc@1: {np.mean(top1_acc):.6f} "
             )
 
+
 def test(args, model, test_loader, device):
     model.eval()
     criterion = nn.CrossEntropyLoss()
diff --git a/examples/ensembling/parallel_train.py b/examples/ensembling/parallel_train.py
index f8a90e6bb..53a310a1c 100644
--- a/examples/ensembling/parallel_train.py
+++ b/examples/ensembling/parallel_train.py
@@ -19,6 +19,8 @@
 DEVICE = 'cpu'
 
 # Step 1: Make some spirals
+
+
 def make_spirals(n_samples, noise_std=0., rotations=1.):
     ts = torch.linspace(0, 1, n_samples, device=DEVICE)
     rs = ts ** 0.5
@@ -31,6 +33,7 @@ def make_spirals(n_samples, noise_std=0., rotations=1.):
     points = torch.stack([xs, ys], dim=1)
     return points, labels
 
+
 points, labels = make_spirals(100, noise_std=0.05)
 
 
@@ -51,6 +54,7 @@ def forward(self, x):
         x = F.log_softmax(x, -1)
         return x
 
+
 loss_fn = nn.NLLLoss()
 
 # Step 3: Make the model functional(!!) and define a training function.
@@ -58,6 +62,7 @@ def forward(self, x):
 # https://github.com/pytorch/pytorch/issues/49171
 func_model, weights = make_functional(MLPClassifier().to(DEVICE))
 
+
 def train_step_fn(weights, batch, targets, lr=0.2):
     def compute_loss(weights, batch, targets):
         output = func_model(weights, batch)
@@ -85,10 +90,13 @@ def step4():
         if i % 100 == 0:
             print(loss)
 
+
 step4()
 
 # Step 5: We're ready for multiple models. Let's define an init_fn
 # that, given a number of models, returns to us all of the weights.
+
+
 def init_fn(num_models):
     models = [MLPClassifier() for _ in range(num_models)]
     _, params, _ = combine_state_for_ensemble(models)
@@ -97,6 +105,8 @@ def init_fn(num_models):
 # Step 6: Now, can we try multiple models at the same time?
 # The answer is: yes! `loss` is a 2-tuple, and we can see that the value keeps
 # on decreasing
+
+
 def step6():
     parallel_train_step_fn = vmap(train_step_fn, in_dims=(0, None, None))
     batched_weights = init_fn(num_models=2)
@@ -105,6 +115,7 @@ def step6():
         if i % 200 == 0:
             print(loss)
 
+
 step6()
 
 # Step 7: Now, the flaw with step 6 is that we were training on the same exact
diff --git a/examples/lennard_jones/lennard_jones.py b/examples/lennard_jones/lennard_jones.py
index bcfc9762f..24df3098a 100644
--- a/examples/lennard_jones/lennard_jones.py
+++ b/examples/lennard_jones/lennard_jones.py
@@ -10,13 +10,16 @@
 sigma = 0.5
 epsilon = 4.
 
+
 def lennard_jones(r):
     return epsilon * ((sigma / r)**12 - (sigma / r)**6)
 
+
 def lennard_jones_force(r):
     """Get magnitude of LJ force"""
     return -epsilon * ((-12 * sigma**12 / r**13) + (6 * sigma**6 / r**7))
 
+
 training_size = 1000
 r = torch.linspace(0.5, 2 * sigma, requires_grad=True)
 
@@ -40,6 +43,7 @@ def lennard_jones_force(r):
     nn.Linear(16, 1)
 )
 
+
 def make_prediction(model, drs):
     norms = torch.norm(drs, dim=1).reshape(-1, 1)
     energies = model(norms)
@@ -48,9 +52,11 @@ def make_prediction(model, drs):
     forces = -network_derivs * drs / norms
     return energies, forces
 
+
 def loss_fn(energies, forces, predicted_energies, predicted_forces):
     return mse_loss(energies, predicted_energies) + 0.01 * mse_loss(forces, predicted_forces) / 3
 
+
 optimiser = torch.optim.Adam(model.parameters(), lr=1e-3)
 
 for epoch in range(400):
diff --git a/examples/maml_omniglot/maml-omniglot-higher.py b/examples/maml_omniglot/maml-omniglot-higher.py
index c5ad2b68d..8f6e017f2 100755
--- a/examples/maml_omniglot/maml-omniglot-higher.py
+++ b/examples/maml_omniglot/maml-omniglot-higher.py
@@ -27,26 +27,22 @@
 https://github.com/bamos/HowToTrainYourMAMLPytorch
 """
 
+from support.omniglot_loaders import OmniglotNShot
+import higher
+import torch.optim as optim
+import torch.nn.functional as F
+from torch import nn
+import torch
+import matplotlib.pyplot as plt
 import argparse
 import time
-import typing
 
 import pandas as pd
 import numpy as np
 import matplotlib as mpl
 mpl.use('Agg')
-import matplotlib.pyplot as plt
 plt.style.use('bmh')
 
-import torch
-from torch import nn
-import torch.nn.functional as F
-import torch.optim as optim
-
-import higher
-
-from support.omniglot_loaders import OmniglotNShot
-
 
 def main():
     argparser = argparse.ArgumentParser()
@@ -199,12 +195,10 @@ def test(db, net, device, epoch, log):
     qry_losses = []
     qry_accs = []
 
-    for batch_idx in range(n_test_iter):
+    for _ in range(n_test_iter):
         x_spt, y_spt, x_qry, y_qry = db.next('test')
 
-
         task_num, setsz, c_, h, w = x_spt.size()
-        querysz = x_qry.size(1)
 
         # TODO: Maybe pull this out into a separate module so it
         # doesn't have to be duplicated between `train` and `test`?
@@ -243,8 +237,6 @@ def test(db, net, device, epoch, log):
     })
 
 
-
-
 def plot(log):
     # Generally you should pull your plotting code out of your training
     # script but we are doing it here for brevity.
diff --git a/examples/maml_omniglot/maml-omniglot-ptonly.py b/examples/maml_omniglot/maml-omniglot-ptonly.py
index 6f786ae33..594237ee7 100755
--- a/examples/maml_omniglot/maml-omniglot-ptonly.py
+++ b/examples/maml_omniglot/maml-omniglot-ptonly.py
@@ -27,27 +27,22 @@
 https://github.com/bamos/HowToTrainYourMAMLPytorch
 """
 
+from support.omniglot_loaders import OmniglotNShot
+from functorch import make_functional_with_buffers
+import torch.optim as optim
+import torch.nn.functional as F
+from torch import nn
+import torch
+import matplotlib.pyplot as plt
 import argparse
 import time
-import typing
 
 import pandas as pd
 import numpy as np
 import matplotlib as mpl
 mpl.use('Agg')
-import matplotlib.pyplot as plt
 plt.style.use('bmh')
 
-import torch
-from torch import nn
-import torch.nn.functional as F
-import torch.optim as optim
-from functorch import make_functional_with_buffers, vmap, grad
-
-import higher
-
-from support.omniglot_loaders import OmniglotNShot
-
 
 def main():
     argparser = argparse.ArgumentParser()
@@ -236,8 +231,6 @@ def test(db, net, device, epoch, log):
     })
 
 
-
-
 def plot(log):
     # Generally you should pull your plotting code out of your training
     # script but we are doing it here for brevity.
diff --git a/examples/maml_omniglot/maml-omniglot-transforms.py b/examples/maml_omniglot/maml-omniglot-transforms.py
index a8a397d28..25b0c3fda 100755
--- a/examples/maml_omniglot/maml-omniglot-transforms.py
+++ b/examples/maml_omniglot/maml-omniglot-transforms.py
@@ -27,33 +27,28 @@
 https://github.com/bamos/HowToTrainYourMAMLPytorch
 """
 
+from support.omniglot_loaders import OmniglotNShot
+from functorch import make_functional_with_buffers, vmap, grad
+import functorch
+import torch.optim as optim
+import torch.nn.functional as F
+from torch import nn
+import torch
+import matplotlib.pyplot as plt
 import argparse
 import time
-import typing
 import functools
 
 import pandas as pd
 import numpy as np
 import matplotlib as mpl
 mpl.use('Agg')
-import matplotlib.pyplot as plt
 plt.style.use('bmh')
 
-import torch
-from torch import nn
-import torch.nn.functional as F
-import torch.optim as optim
-
-import functorch
-from functorch import make_functional_with_buffers, vmap, grad
 
 # Squash the warning spam
 functorch._C._set_vmap_fallback_warning_enabled(False)
 
-import higher
-
-from support.omniglot_loaders import OmniglotNShot
-
 
 def main():
     argparser = argparse.ArgumentParser()
@@ -245,8 +240,6 @@ def test(db, net, device, epoch, log):
     })
 
 
-
-
 def plot(log):
     # Generally you should pull your plotting code out of your training
     # script but we are doing it here for brevity.
diff --git a/examples/maml_omniglot/support/omniglot_loaders.py b/examples/maml_omniglot/support/omniglot_loaders.py
index 31118a0b4..b712b9b31 100644
--- a/examples/maml_omniglot/support/omniglot_loaders.py
+++ b/examples/maml_omniglot/support/omniglot_loaders.py
@@ -17,15 +17,15 @@
 #     https://github.com/dragen1860/MAML-Pytorch/blob/master/omniglot.py
 #     https://github.com/dragen1860/MAML-Pytorch/blob/master/omniglotNShot.py
 
-import  torchvision.transforms as transforms
-from    PIL import Image
-import  numpy as np
+import torchvision.transforms as transforms
+from PIL import Image
+import numpy as np
 
 import torch
-import  torch.utils.data as data
-import  os
-import  os.path
-import  errno
+import torch.utils.data as data
+import os
+import os.path
+import errno
 
 
 class Omniglot(data.Dataset):
@@ -79,7 +79,7 @@ def __len__(self):
 
     def _check_exists(self):
         return os.path.exists(os.path.join(self.root, self.processed_folder, "images_evaluation")) and \
-               os.path.exists(os.path.join(self.root, self.processed_folder, "images_background"))
+            os.path.exists(os.path.join(self.root, self.processed_folder, "images_background"))
 
     def download(self):
         from six.moves import urllib
@@ -158,7 +158,7 @@ def __init__(self, root, batchsz, n_way, k_shot, k_query, imgsz, device=None):
                      lambda x: x.resize((imgsz, imgsz)),
                      lambda x: np.reshape(x, (imgsz, imgsz, 1)),
                      lambda x: np.transpose(x, [2, 0, 1]),
-                     lambda x: x/255.]),
+                     lambda x: x / 255.]),
             )
 
             temp = dict()  # {label:img1, img2..., 20 imgs, label2: img1, img2,... in total, 1623 label}
@@ -196,7 +196,7 @@ def __init__(self, root, batchsz, n_way, k_shot, k_query, imgsz, device=None):
         self.n_way = n_way  # n way
         self.k_shot = k_shot  # k shot
         self.k_query = k_query  # k query
-        assert (k_shot + k_query) <=20
+        assert (k_shot + k_query) <= 20
 
         # save pointer of current read batch in total cache
         self.indexes = {"train": 0, "test": 0}
@@ -269,7 +269,6 @@ def load_data_cache(self, data_pack):
                 x_qrys.append(x_qry)
                 y_qrys.append(y_qry)
 
-
             # [b, setsz, 1, 84, 84]
             x_spts = np.array(x_spts).astype(np.float32).reshape(self.batchsz, setsz, 1, self.resize, self.resize)
             y_spts = np.array(y_spts).astype(np.int).reshape(self.batchsz, setsz)
diff --git a/examples/maml_regression/evjang.py b/examples/maml_regression/evjang.py
index b64030f4f..c0e590d25 100644
--- a/examples/maml_regression/evjang.py
+++ b/examples/maml_regression/evjang.py
@@ -10,15 +10,14 @@
 #       http://www.apache.org/licenses/LICENSE-2.0
 #
 
+import matplotlib.pyplot as plt
 import math
-import random
 import torch
 import numpy as np
-from torch import nn
 from torch.nn import functional as F
 import matplotlib as mpl
 mpl.use('Agg')
-import matplotlib.pyplot as plt
+
 
 def net(x, params):
     x = F.linear(x, params[0], params[1])
@@ -30,14 +29,15 @@ def net(x, params):
     x = F.linear(x, params[4], params[5])
     return x
 
+
 params = [
     torch.Tensor(40, 1).uniform_(-1., 1.).requires_grad_(),
     torch.Tensor(40).zero_().requires_grad_(),
 
-    torch.Tensor(40, 40).uniform_(-1./math.sqrt(40), 1./math.sqrt(40)).requires_grad_(),
+    torch.Tensor(40, 40).uniform_(-1. / math.sqrt(40), 1. / math.sqrt(40)).requires_grad_(),
     torch.Tensor(40).zero_().requires_grad_(),
 
-    torch.Tensor(1, 40).uniform_(-1./math.sqrt(40), 1./math.sqrt(40)).requires_grad_(),
+    torch.Tensor(1, 40).uniform_(-1. / math.sqrt(40), 1. / math.sqrt(40)).requires_grad_(),
     torch.Tensor(1).zero_().requires_grad_(),
 ]
 
@@ -47,6 +47,8 @@ def net(x, params):
 K = 20
 losses = []
 num_tasks = 4
+
+
 def sample_tasks(outer_batch_size, inner_batch_size):
     # Select amplitude and phase for the task
     As = []
@@ -54,6 +56,7 @@ def sample_tasks(outer_batch_size, inner_batch_size):
     for _ in range(outer_batch_size):
         As.append(np.random.uniform(low=0.1, high=.5))
         phases.append(np.random.uniform(low=0., high=np.pi))
+
     def get_batch():
         xs, ys = [], []
         for A, phase in zip(As, phases):
@@ -66,9 +69,11 @@ def get_batch():
     x2, y2 = get_batch()
     return x1, y1, x2, y2
 
+
 for it in range(20000):
     loss2 = 0.0
     opt.zero_grad()
+
     def get_loss_for_task(x1, y1, x2, y2):
         f = net(x1, params)
         loss = F.mse_loss(f, y1)
@@ -77,14 +82,14 @@ def get_loss_for_task(x1, y1, x2, y2):
         # We want to differentiate through the SGD update steps and get higher order
         # derivatives in the backward pass.
         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(x2, new_params)
         return F.mse_loss(v_f, y2)
 
     task = sample_tasks(num_tasks, K)
     inner_losses = [get_loss_for_task(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)
     loss2.backward()
 
     opt.step()
@@ -97,7 +102,7 @@ def get_loss_for_task(x1, y1, x2, y2):
 t_b = torch.tensor(0.0).uniform_(0.0, math.pi)
 
 t_x = torch.empty(4, 1).uniform_(-5, 5)
-t_y = t_A*torch.sin(t_x + t_b)
+t_y = t_A * torch.sin(t_x + t_b)
 
 opt.zero_grad()
 
@@ -107,11 +112,11 @@ def get_loss_for_task(x1, y1, x2, y2):
     t_loss = F.l1_loss(t_f, t_y)
 
     grads = torch.autograd.grad(t_loss, t_params, create_graph=True)
-    t_params = [(t_params[i] - alpha*grads[i]) for i in range(len(params))]
+    t_params = [(t_params[i] - alpha * grads[i]) for i in range(len(params))]
 
 
-test_x = torch.arange(-2*math.pi, 2*math.pi, step=0.01).unsqueeze(1)
-test_y = t_A*torch.sin(test_x + t_b)
+test_x = torch.arange(-2 * math.pi, 2 * math.pi, step=0.01).unsqueeze(1)
+test_y = t_A * torch.sin(test_x + t_b)
 
 test_f = net(test_x, t_params)
 
@@ -121,5 +126,5 @@ def get_loss_for_task(x1, y1, x2, y2):
 plt.legend()
 plt.savefig('maml-sine.png')
 plt.figure()
-plt.plot(np.convolve(losses, [.05]*20))
+plt.plot(np.convolve(losses, [.05] * 20))
 plt.savefig('losses.png')
diff --git a/examples/maml_regression/evjang_transforms.py b/examples/maml_regression/evjang_transforms.py
index b2206fa02..5d0dafdbd 100644
--- a/examples/maml_regression/evjang_transforms.py
+++ b/examples/maml_regression/evjang_transforms.py
@@ -10,17 +10,15 @@
 #       http://www.apache.org/licenses/LICENSE-2.0
 #
 
+from functorch import grad, vmap
+import matplotlib.pyplot as plt
 import math
-import random
 import torch
 import numpy as np
-from torch import nn
 from torch.nn import functional as F
 import matplotlib as mpl
 mpl.use('Agg')
-import matplotlib.pyplot as plt
 
-from functorch import grad, vmap
 
 def net(params, x):
     x = F.linear(x, params[0], params[1])
@@ -32,27 +30,33 @@ def net(params, x):
     x = F.linear(x, params[4], params[5])
     return x
 
+
 params = [
     torch.Tensor(40, 1).uniform_(-1., 1.).requires_grad_(),
     torch.Tensor(40).zero_().requires_grad_(),
 
-    torch.Tensor(40, 40).uniform_(-1./math.sqrt(40), 1./math.sqrt(40)).requires_grad_(),
+    torch.Tensor(40, 40).uniform_(-1. / math.sqrt(40), 1. / math.sqrt(40)).requires_grad_(),
     torch.Tensor(40).zero_().requires_grad_(),
 
-    torch.Tensor(1, 40).uniform_(-1./math.sqrt(40), 1./math.sqrt(40)).requires_grad_(),
+    torch.Tensor(1, 40).uniform_(-1. / math.sqrt(40), 1. / math.sqrt(40)).requires_grad_(),
     torch.Tensor(1).zero_().requires_grad_(),
 ]
 
 # The prototype doesn't like F.mse_loss.
+
+
 def mse_loss(x, y):
     return torch.mean((x - y) ** 2)
 
+
 opt = torch.optim.Adam(params, lr=1e-3)
 alpha = 0.1
 
 K = 20
 losses = []
 num_tasks = 4
+
+
 def sample_tasks(outer_batch_size, inner_batch_size):
     # Select amplitude and phase for the task
     As = []
@@ -60,6 +64,7 @@ def sample_tasks(outer_batch_size, inner_batch_size):
     for _ in range(outer_batch_size):
         As.append(np.random.uniform(low=0.1, high=.5))
         phases.append(np.random.uniform(low=0., high=np.pi))
+
     def get_batch():
         xs, ys = [], []
         for A, phase in zip(As, phases):
@@ -72,9 +77,11 @@ def get_batch():
     x2, y2 = get_batch()
     return x1, y1, x2, y2
 
+
 for it in range(20000):
     loss2 = 0.0
     opt.zero_grad()
+
     def get_loss_for_task(x1, y1, x2, y2):
         def inner_loss(params, x1, y1):
             f = net(params, x1)
@@ -82,14 +89,14 @@ def inner_loss(params, x1, y1):
             return loss
 
         grads = grad(inner_loss)(tuple(params), x1, y1)
-        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)
 
     task = sample_tasks(num_tasks, K)
     inner_losses = vmap(get_loss_for_task)(task[0], task[1], task[2], task[3])
-    loss2 = sum(inner_losses)/len(inner_losses)
+    loss2 = sum(inner_losses) / len(inner_losses)
     loss2.backward()
 
     opt.step()
@@ -102,7 +109,7 @@ def inner_loss(params, x1, y1):
 t_b = torch.tensor(0.0).uniform_(0.0, math.pi)
 
 t_x = torch.empty(4, 1).uniform_(-5, 5)
-t_y = t_A*torch.sin(t_x + t_b)
+t_y = t_A * torch.sin(t_x + t_b)
 
 opt.zero_grad()
 
@@ -112,11 +119,11 @@ def inner_loss(params, x1, y1):
     t_loss = F.l1_loss(t_f, t_y)
 
     grads = torch.autograd.grad(t_loss, t_params, create_graph=True)
-    t_params = [(t_params[i] - alpha*grads[i]) for i in range(len(params))]
+    t_params = [(t_params[i] - alpha * grads[i]) for i in range(len(params))]
 
 
-test_x = torch.arange(-2*math.pi, 2*math.pi, step=0.01).unsqueeze(1)
-test_y = t_A*torch.sin(test_x + t_b)
+test_x = torch.arange(-2 * math.pi, 2 * math.pi, step=0.01).unsqueeze(1)
+test_y = t_A * torch.sin(test_x + t_b)
 
 test_f = net(test_x, t_params)
 
@@ -126,5 +133,5 @@ def inner_loss(params, x1, y1):
 plt.legend()
 plt.savefig('maml-sine.png')
 plt.figure()
-plt.plot(np.convolve(losses, [.05]*20))
+plt.plot(np.convolve(losses, [.05] * 20))
 plt.savefig('losses.png')
diff --git a/examples/maml_regression/evjang_transforms_module.py b/examples/maml_regression/evjang_transforms_module.py
index acca8afad..732b4f5fa 100644
--- a/examples/maml_regression/evjang_transforms_module.py
+++ b/examples/maml_regression/evjang_transforms_module.py
@@ -10,17 +10,16 @@
 #       http://www.apache.org/licenses/LICENSE-2.0
 #
 
+from functorch import grad, vmap, make_functional
+import matplotlib.pyplot as plt
 import math
-import random
 import torch
 import numpy as np
 from torch import nn
 from torch.nn import functional as F
 import matplotlib as mpl
 mpl.use('Agg')
-import matplotlib.pyplot as plt
 
-from functorch import grad, vmap, make_functional
 
 class ThreeLayerNet(nn.Module):
     def __init__(self):
@@ -40,9 +39,12 @@ def forward(self, x):
         return x
 
 # The prototype doesn't like F.mse_loss.
+
+
 def mse_loss(x, y):
     return torch.mean((x - y) ** 2)
 
+
 net, params = make_functional(ThreeLayerNet())
 opt = torch.optim.Adam(params, lr=1e-3)
 alpha = 0.1
@@ -50,6 +52,8 @@ def mse_loss(x, y):
 K = 20
 losses = []
 num_tasks = 4
+
+
 def sample_tasks(outer_batch_size, inner_batch_size):
     # Select amplitude and phase for the task
     As = []
@@ -57,6 +61,7 @@ def sample_tasks(outer_batch_size, inner_batch_size):
     for _ in range(outer_batch_size):
         As.append(np.random.uniform(low=0.1, high=.5))
         phases.append(np.random.uniform(low=0., high=np.pi))
+
     def get_batch():
         xs, ys = [], []
         for A, phase in zip(As, phases):
@@ -69,9 +74,11 @@ def get_batch():
     x2, y2 = get_batch()
     return x1, y1, x2, y2
 
+
 for it in range(20000):
     loss2 = 0.0
     opt.zero_grad()
+
     def get_loss_for_task(x1, y1, x2, y2):
         def inner_loss(params, x1, y1):
             f = net(params, x1)
@@ -79,14 +86,14 @@ def inner_loss(params, x1, y1):
             return loss
 
         grads = grad(inner_loss)(params, x1, y1)
-        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)
 
     task = sample_tasks(num_tasks, K)
     inner_losses = vmap(get_loss_for_task)(task[0], task[1], task[2], task[3])
-    loss2 = sum(inner_losses)/len(inner_losses)
+    loss2 = sum(inner_losses) / len(inner_losses)
     loss2.backward()
 
     opt.step()
@@ -99,7 +106,7 @@ def inner_loss(params, x1, y1):
 t_b = torch.tensor(0.0).uniform_(0.0, math.pi)
 
 t_x = torch.empty(4, 1).uniform_(-5, 5)
-t_y = t_A*torch.sin(t_x + t_b)
+t_y = t_A * torch.sin(t_x + t_b)
 
 opt.zero_grad()
 
@@ -109,11 +116,11 @@ def inner_loss(params, x1, y1):
     t_loss = F.l1_loss(t_f, t_y)
 
     grads = torch.autograd.grad(t_loss, t_params, create_graph=True)
-    t_params = [(t_params[i] - alpha*grads[i]) for i in range(len(params))]
+    t_params = [(t_params[i] - alpha * grads[i]) for i in range(len(params))]
 
 
-test_x = torch.arange(-2*math.pi, 2*math.pi, step=0.01).unsqueeze(1)
-test_y = t_A*torch.sin(test_x + t_b)
+test_x = torch.arange(-2 * math.pi, 2 * math.pi, step=0.01).unsqueeze(1)
+test_y = t_A * torch.sin(test_x + t_b)
 
 test_f = net(t_params, test_x)
 
@@ -123,5 +130,5 @@ def inner_loss(params, x1, y1):
 plt.legend()
 plt.savefig('maml-sine.png')
 plt.figure()
-plt.plot(np.convolve(losses, [.05]*20))
+plt.plot(np.convolve(losses, [.05] * 20))
 plt.savefig('losses.png')
diff --git a/functorch/__init__.py b/functorch/__init__.py
index 37e77da82..0888374ff 100644
--- a/functorch/__init__.py
+++ b/functorch/__init__.py
@@ -41,8 +41,8 @@ def cross_entropy(input, target, weight=None, size_average=None,
         target = target.unsqueeze(0)
 
     result = _old_cross_entropy(
-            input, target, weight, size_average,
-            ignore_index, reduce, reduction, **kwargs)
+        input, target, weight, size_average,
+        ignore_index, reduce, reduction, **kwargs)
     if reduction == 'none':
         return result.squeeze(0)
     return result
diff --git a/functorch/_src/aot_autograd.py b/functorch/_src/aot_autograd.py
index 976a96633..ebfdc8907 100644
--- a/functorch/_src/aot_autograd.py
+++ b/functorch/_src/aot_autograd.py
@@ -1,5 +1,4 @@
-from functorch import make_fx
-import time
+import os
 import torch
 import torch.nn as nn
 from functorch import make_functional_with_buffers, make_fx
@@ -10,20 +9,18 @@
 import torch.utils._pytree as pytree
 import torch.utils.dlpack
 from torch.fx.passes import graph_drawer
-import operator
-import os
-import inspect
 from functorch._C import CompileCache
-from functools import partial
-from typing import Callable
 from .python_key import pythonkey_decompose
 from .decompositions import register_decomposition
 
-pytree._register_pytree_node(immutable_collections.immutable_list, lambda x: (list(x), None), lambda x, c: immutable_collections.immutable_list(x))
-pytree._register_pytree_node(immutable_collections.immutable_dict, lambda x: (list(x.values()), list(x.keys())), lambda x, c: immutable_collections.immutable_dict({key: value for key, value in zip(c, x)}))
+pytree._register_pytree_node(immutable_collections.immutable_list, lambda x: (
+    list(x), None), lambda x, c: immutable_collections.immutable_list(x))
+pytree._register_pytree_node(immutable_collections.immutable_dict, lambda x: (list(x.values()), list(
+    x.keys())), lambda x, c: immutable_collections.immutable_dict({key: value for key, value in zip(c, x)}))
 
 aten = torch.ops.aten
 
+
 def draw_graph(traced: torch.fx.GraphModule, fname: str, figname: str = "fx_graph"):
     base, ext = os.path.splitext(fname)
     if not ext:
@@ -34,6 +31,8 @@ def draw_graph(traced: torch.fx.GraphModule, fname: str, figname: str = "fx_grap
     getattr(x, "write_" + ext.lstrip("."))(f"{base}{ext}")
 
 # todo(chilli): clean this up/make it more understandable
+
+
 def default_partition(fx_module: fx.GraphModule, _joint_inputs):
     bw_nodes = set()
     saved_nodes = set()
@@ -72,7 +71,7 @@ def add_saved(a):
 
     for node in fx_module.graph.nodes:
         if node in bw_nodes or node in bw_outputs:
-            value_remap[node] = bw_graph.node_copy(node, lambda n : value_remap[n])
+            value_remap[node] = bw_graph.node_copy(node, lambda n: value_remap[n])
 
     assert(num_fwd_outputs + num_bwd_outputs == len(output_node.args[0]))
     bwd_outputs = [value_remap[i] for i in bw_outputs]
@@ -85,9 +84,10 @@ def add_saved(a):
     value_remap = {}
     for node in fx_module.graph.nodes:
         if node not in bw_nodes and node.op != 'output':
-            value_remap[node] = fw_graph.node_copy(node, lambda n : value_remap[n])
+            value_remap[node] = fw_graph.node_copy(node, lambda n: value_remap[n])
 
-    fwd_outputs = [value_remap[i] for i in output_node.args[0][:num_fwd_outputs]] + [value_remap[n] for n in saved_nodes]
+    fwd_outputs = [value_remap[i] for i in output_node.args[0]
+                   [:num_fwd_outputs]] + [value_remap[n] for n in saved_nodes]
     if len(fwd_outputs) == 1:
         fwd_outputs = fwd_outputs[0]
     fw_graph.output(fwd_outputs)
@@ -96,11 +96,15 @@ def add_saved(a):
     bw_module.graph.lint()
     return fw_module, bw_module
 
+
 class InvalidNodeBase(object):
     def __repr__(self):
         return "Invalid Node"
+
+
 InvalidNode = InvalidNodeBase()
 
+
 def _extract_graph_with_inputs_outputs(joint_graph, inputs, outputs):
     """
     Given a graph, extracts out a subgraph that takes the specified nodes as inputs and returns the specified outputs.
@@ -153,6 +157,7 @@ def map_arg_to_proxy(x):
     new_graph.lint()
     return new_graph
 
+
 def partition_with_recompute_fwd_in_bwd(joint_module: fx.GraphModule, _joint_inputs):
     """
     Partitions the joint graph such that the backward recomputes the forward.
@@ -163,7 +168,6 @@ def partition_with_recompute_fwd_in_bwd(joint_module: fx.GraphModule, _joint_inp
     resulting graphs through dead code elimintation.
     """
 
-
     def is_primal(node):
         return node.op == "placeholder" and "tangents" not in node.target
 
@@ -198,7 +202,7 @@ def is_tangent(node):
                     break
 
     # Now, we re-generate the fwd/bwd graphs.
-    # NB: This might increase compilation time, but I doubt it matters  
+    # NB: This might increase compilation time, but I doubt it matters
     fwd_graph = _extract_graph_with_inputs_outputs(joint_module.graph, primal_inputs, fwd_outputs + saved_values)
     bwd_graph = _extract_graph_with_inputs_outputs(joint_module.graph, saved_values + tangent_inputs, bwd_outputs)
 
@@ -207,20 +211,23 @@ def is_tangent(node):
 
     return fwd_module, bwd_module
 
+
 def create_joint_forward_backward(fn):
     def joint_forward_backward(primals, tangents):
         out = fn(*primals)
         primals = [p for p in pytree.tree_flatten(primals)[0] if p.requires_grad]
         backward_out = []
-        if primals: # todo(chilli): Make it support it if not all outputs have gradients
+        if primals:  # todo(chilli): Make it support it if not all outputs have gradients
             backward_out = torch.autograd.grad(out, primals, grad_outputs=tangents, allow_unused=True)
         return out, backward_out
     return joint_forward_backward
 
+
 def draw_joint_graph(graph, joint_inputs, file_name="full_graph.png"):
     draw_graph(graph, file_name)
     return default_partition(graph, joint_inputs)
 
+
 def normalize_as_list(x):
     if isinstance(x, tuple):
         return list(x)
@@ -228,6 +235,7 @@ def normalize_as_list(x):
         return x
     return [x]
 
+
 def create_compiled_function(flat_fn, fw_compiler, bw_compiler, partition_fn, decompose):
 
     # putting these decompositions here since they shouldn't always be used
@@ -302,11 +310,13 @@ class _CompileCache(CompileCache):
 compile_cache = None
 
 # Inspired by autodidax (thanks!)
+
+
 class PytreeThunk:
     spec = None
     # These are some kinda dumb microoptimizations that save about 3-4 us of overhead.
-    is_simple = None # if the output spec is a tuple/list, we won't bother unflattening it.
-    is_really_simple = None # if the output spec is a LeafSpec
+    is_simple = None  # if the output spec is a tuple/list, we won't bother unflattening it.
+    is_really_simple = None  # if the output spec is a LeafSpec
 
     def set(self, spec):
         assert self.spec is None or self.spec == spec
@@ -323,6 +333,7 @@ def unflatten(self, x):
             return x
         return pytree.tree_unflatten(x, self.spec)
 
+
 def compiled_function(
     fn, fw_compiler, bw_compiler, partition_fn=default_partition, decompose=False, hasher_type="StaticShapeHasher"
 ):
@@ -349,6 +360,7 @@ def returned_function(*args, **kwargs):
             # Compile a new function
             flattened_args, args_spec = pytree.tree_flatten((args, kwargs))
             out_spec = PytreeThunk()
+
             def flat_fn(*args):
                 nonlocal out_spec
                 args, kwargs = pytree.tree_unflatten(args, args_spec)
@@ -405,5 +417,6 @@ def forward(self, *args, **kwargs):
 
     return CompiledModule()
 
+
 aot_function = compiled_function
 aot_module = compiled_module
diff --git a/functorch/_src/compilers.py b/functorch/_src/compilers.py
index 31169719e..6e17c511f 100644
--- a/functorch/_src/compilers.py
+++ b/functorch/_src/compilers.py
@@ -3,6 +3,7 @@
 from .aot_autograd import draw_graph
 import time
 
+
 def ts_compile(fx_g, _):
     for node in fx_g.graph.nodes:
         if node.target == torch.ops.aten.new_zeros:
@@ -44,14 +45,17 @@ def ts_compile(fx_g, _):
     f = torch.jit.optimize_for_inference(f)
     return f
 
+
 def _draw_graph_compile(fx_g, _, name):
     draw_graph(fx_g, name)
     return fx_g
 
+
 def draw_graph_compile(name):
     return partial(draw_graph_compile, name=name)
 
-def _tvm_compile(fx_module, example_inputs, name = None):
+
+def _tvm_compile(fx_module, example_inputs, name=None):
     import tvm
     from tvm import relay, auto_scheduler
     from tvm.contrib import graph_executor
@@ -85,8 +89,8 @@ def _tvm_compile(fx_module, example_inputs, name = None):
     with auto_scheduler.ApplyHistoryBest(log_file):
         with tvm.transform.PassContext(opt_level=3, config={"relay.backend.use_auto_scheduler": True}):
             lib = relay.build(mod, target=target, params=params)
-    dtype = "float32"
     m = graph_executor.GraphModule(lib["default"](dev))
+
     def exec_tvm(*args):
         for idx, arg in enumerate(args, 0):
             if arg.dim() != 0:
@@ -97,9 +101,11 @@ def exec_tvm(*args):
         return outs
     return exec_tvm
 
+
 def tvm_compile(name):
     return partial(tvm_compile, name=name)
 
+
 def nop(f, _):
     print(f.code)
-    return f
\ No newline at end of file
+    return f
diff --git a/functorch/_src/custom_function.py b/functorch/_src/custom_function.py
index 70635c7c6..028a246c6 100644
--- a/functorch/_src/custom_function.py
+++ b/functorch/_src/custom_function.py
@@ -3,6 +3,7 @@
 
 m = functorch._C._dispatch_library("FRAGMENT", "aten", "")
 
+
 def custom_vjp(name, filter_fn, fwd_fn, bwd_fn):
     m.def_(f"{name}(Tensor[] args) -> Tensor[]")
     m.impl(f"{name}", "CompositeImplicitAutograd", fwd_fn)
@@ -11,8 +12,8 @@ def custom_vjp(name, filter_fn, fwd_fn, bwd_fn):
     m.impl(f"{name}_vjp", "CompositeImplicitAutograd", bwd_fn)
 
     # TODO: it looks like the autograd alias key doesn't work
-    m.gen_backward_binding(f"{name}", "AutogradCPU");
-    m.gen_backward_binding(f"{name}", "AutogradCUDA");
+    m.gen_backward_binding(f"{name}", "AutogradCPU")
+    m.gen_backward_binding(f"{name}", "AutogradCUDA")
 
     def wrapped(*args):
         return filter_fn(getattr(torch.ops.aten, name)(args))
diff --git a/functorch/_src/decompositions.py b/functorch/_src/decompositions.py
index 800b8bd49..7d9fd4e69 100644
--- a/functorch/_src/decompositions.py
+++ b/functorch/_src/decompositions.py
@@ -1,67 +1,89 @@
 import torch
 from torch import Tensor
-from typing import Any, Dict, NamedTuple, Optional, Set, Tuple, List, Callable, Union
+from typing import Optional, List
 from enum import Enum
 
 aten = torch.ops.aten
 
 decomposition_table = {}
 
+
 def register_decomposition(aten_op):
     def decomposition_decorator(f):
         decomposition_table[aten_op] = f
         return f
     return decomposition_decorator
 
+
 class Reduction(Enum):
     NONE = 0
     MEAN = 1
     SUM = 2
 
+
 @register_decomposition(aten.tanh_backward)
 def tanh_backward_decomposition(out_grad: Tensor, y: Tensor):
     return out_grad * (1 - y * y)
 
+
 @register_decomposition(aten.sigmoid_backward)
 def sigmoid_backward_decomposition(out_grad: Tensor, y: Tensor):
     return out_grad * (y * (1 - y))
 
+
 @register_decomposition(aten.softplus_backward)
 # The out argument seems to always be ignored?
 def softplus_backward_decomposition(out_grad: Tensor, x: Tensor, beta: float, threshold: float, out):
     z = (x * beta).exp()
     return aten.where((x * beta) > threshold, out_grad, out_grad * z / (z + 1.0))
 
+
 @register_decomposition(aten.elu_backward)
-def elu_backward_decomposition(grad_output: Tensor, alpha: float, scale: float, input_scale: float, is_result: bool, self_or_result: Tensor):
+def elu_backward_decomposition(
+    grad_output: Tensor, alpha: float, scale: float, input_scale: float, is_result: bool, self_or_result: Tensor
+):
     negcoef = alpha * scale
     poscoef = scale
     negiptcoef = input_scale
     if is_result:
-        return aten.where(self_or_result <= 0, grad_output * negiptcoef * (self_or_result + negcoef), self_or_result * poscoef)
+        return aten.where(
+            self_or_result <= 0,
+            grad_output * negiptcoef * (self_or_result + negcoef),
+            self_or_result * poscoef
+        )
     else:
-        return aten.where(self_or_result <= 0, grad_output * negiptcoef * negcoef * aten.exp(self_or_result * negiptcoef), grad_output * poscoef)
+        return aten.where(
+            self_or_result <= 0,
+            grad_output * negiptcoef * negcoef * aten.exp(self_or_result * negiptcoef),
+            grad_output * poscoef
+        )
+
 
 @register_decomposition(aten.hardsigmoid_backward)
 def hardsigmoid_backward_decomposition(grad_output: Tensor, self: Tensor):
-    return aten.where((self > -3.0) & (self < 3.0), grad_output * (1.0/6.0), aten.new_zeros(grad_output, ()))
+    return aten.where((self > -3.0) & (self < 3.0), grad_output * (1.0 / 6.0), aten.new_zeros(grad_output, ()))
+
 
 @register_decomposition(aten.hardtanh_backward)
 def hardtanh_backward_decomposition(grad_output: Tensor, self: Tensor, min_val: float, max_val: float):
     return aten.where((self <= min_val) | (self >= max_val), aten.new_zeros(grad_output, ()), grad_output)
 
+
 @register_decomposition(aten.hardshrink_backward)
 def hardshrink_backward(grad_out: Tensor, self: Tensor, lambd: float):
     return aten.where((self >= -lambd) & (self <= lambd), aten.new_zeros(grad_out, ()), grad_out)
 
+
 @register_decomposition(aten.threshold_backward)
 def threshold_backward_decomposition(grad_output: Tensor, self: Tensor, threshold: float):
     return aten.where(self <= threshold, aten.new_zeros(grad_output, ()), grad_output)
 
+
 @register_decomposition(aten.leaky_relu_backward)
 def leaky_relu_backward_decomposition(grad_output: Tensor, self: Tensor, negative_slope: float, self_is_result: bool):
     return aten.where(self > 0, grad_output, grad_output * negative_slope)
 
+
 @register_decomposition(aten.gelu_backward)
 def gelu_backward_decomposition(grad: Tensor, self: Tensor):
     M_SQRT1_2 = 0.70710678118654752440
@@ -72,68 +94,89 @@ def gelu_backward_decomposition(grad: Tensor, self: Tensor):
     pdf = kBeta * aten.exp(self * self * -0.5)
     return grad * (cdf + self * pdf)
 
+
 @register_decomposition(aten.mish_backward)
 def mish_backward_decomposition(grad_output: Tensor, input: Tensor):
-  input_tanh_softplus = aten.tanh(aten.softplus(input))
-  input_sigmoid = aten.sigmoid(input)
-  return grad_output * (input_tanh_softplus + (input * input_sigmoid * (1 - input_tanh_softplus * input_tanh_softplus)))
+    input_tanh_softplus = aten.tanh(aten.softplus(input))
+    input_sigmoid = aten.sigmoid(input)
+    out = (input * input_sigmoid * (1 - input_tanh_softplus * input_tanh_softplus))
+    return grad_output * (input_tanh_softplus + out)
 
 # whyyyy does log_sigmoid do 2 different things for CPU and CUDA >:(
+
+
 @register_decomposition(aten.log_sigmoid_backward)
 def log_sigmoid_backward(grad_output: Tensor, self: Tensor, buffer: Tensor) -> Tensor:
     in_negative = self < 0
     max_deriv = aten.where(in_negative, 1, 0)
     sign = aten.where(in_negative, 1, -1)
-    if grad_output.is_cuda: # buffer is not used on CUDA
+    if grad_output.is_cuda:  # buffer is not used on CUDA
         z = aten.exp(-aten.abs(self))
         return grad_output * (max_deriv - sign * (z / (1 + z)))
     else:
         return (max_deriv - sign * (buffer / (1 + buffer))) * grad_output
 
+
 @register_decomposition(aten.mse_loss_backward)
 def mse_loss_backward(grad_output: Tensor, input: Tensor, target: Tensor, reduction: int):
-    norm = 2./input.numel() if reduction == Reduction.MEAN.value else 2.
+    norm = 2. / input.numel() if reduction == Reduction.MEAN.value else 2.
     return norm * (input - target) * grad_output
 
+
 @register_decomposition(aten.huber_loss_backward)
 def huber_loss_backward(grad_output: Tensor, self: Tensor, target: Tensor, reduction: int, delta: float):
-    norm = 1./self.numel() if reduction == Reduction.MEAN.value else 1.
+    norm = 1. / self.numel() if reduction == Reduction.MEAN.value else 1.
     x = self - target
-    return aten.where(x < -delta, -norm * grad_output * delta, aten.where(x > delta, norm * grad_output * delta, norm * x * grad_output))
+    return aten.where(
+        x < -delta,
+        -norm * grad_output * delta,
+        aten.where(x > delta, norm * grad_output * delta, norm * x * grad_output)
+    )
+
 
 @register_decomposition(aten.slice_backward)
-def slice_backward(grad_output: Tensor, input_sizes: List[int], dim: int, start: int, end: int, step:int):
+def slice_backward(grad_output: Tensor, input_sizes: List[int], dim: int, start: int, end: int, step: int):
     grad_input = aten.new_zeros(grad_output, input_sizes)
     return aten.slice_scatter(grad_input, grad_output, dim, start, end, step)
 
+
 @register_decomposition(aten.select_backward)
 def select_backward(grad_output: Tensor, input_sizes: List[int], dim: int, index: int):
     grad_input = aten.new_zeros(grad_output, input_sizes)
     return aten.select_scatter(grad_input, grad_output, dim, index)
 
+
 @register_decomposition(aten.diagonal_backward)
 def diagonal_backward(grad_output: Tensor, input_sizes: List[int], offset: int, dim1: int, dim2: int):
     grad_input = aten.new_zeros(grad_output, input_sizes)
     return aten.diagonal_scatter(grad_input, grad_output, offset, dim1, dim2)
 
+
 @register_decomposition(aten._softmax_backward_data)
 def _softmax_backward_data(grad_output: Tensor, output: Tensor, dim: int, input_dtype: int):
     new_grad = grad_output * output
     return (new_grad - output * aten.sum(new_grad, dim=dim, keepdim=True))
 
+
 @register_decomposition(aten._log_softmax_backward_data)
 def _log_softmax_backward_data(grad_output: Tensor, output: Tensor, dim: int, input_dtype: int):
     grad_input = grad_output - aten.exp(output) * aten.sum(grad_output, dim=dim, keepdim=True)
     return grad_input
 
+
 @register_decomposition(aten.im2col_backward)
-def im2col_backward(grad_output: Tensor, input_size: List[int], kernel_size: List[int], dilation: List[int], padding: List[int], stride: List[int]) -> Tensor:
+def im2col_backward(
+    grad_output: Tensor, input_size: List[int], kernel_size: List[int],
+    dilation: List[int], padding: List[int], stride: List[int]
+) -> Tensor:
     return aten.col2im(grad_output, input_size, kernel_size, dilation, padding, stride)
 
+
 @register_decomposition(aten.native_dropout_backward)
 def native_dropout_backward(grad_output: Tensor, mask: Tensor, scale: float):
     return grad_output * (mask.type_as(grad_output) * scale)
 
+
 @register_decomposition(aten.logit_backward)
 def logit_backward(grad_output: Tensor, self: Tensor, eps: Optional[float] = None) -> Tensor:
     if eps is not None:
@@ -150,7 +193,6 @@ def logit_backward(grad_output: Tensor, self: Tensor, eps: Optional[float] = Non
             aten.new_full(self, (), float('nan')))
 
 
-
 @register_decomposition(aten.native_dropout)
 def native_dropout_decomposition(input, p, generator=None):
     bool_mask = aten.rand_like(input) < p
@@ -166,4 +208,4 @@ def native_dropout_decomposition(input, p, generator=None):
 
 @register_decomposition(aten._s_where)
 def _s_where_canonicalization(a, b, c):
-    return aten.where(a, b, c)
\ No newline at end of file
+    return aten.where(a, b, c)
diff --git a/functorch/_src/eager_transforms.py b/functorch/_src/eager_transforms.py
index 2d3246eeb..42f1c4d80 100644
--- a/functorch/_src/eager_transforms.py
+++ b/functorch/_src/eager_transforms.py
@@ -32,6 +32,7 @@ def create_differentiable(x):
                          f'got {type(x)}')
     return tree_map(create_differentiable, inps)
 
+
 def _undo_create_differentiable(inps, level=None):
     def unwrap_tensors(x):
         if isinstance(x, torch.Tensor):
@@ -44,29 +45,36 @@ def unwrap_tensors(x):
 
     return tree_map(unwrap_tensors, inps)
 
+
 def _is_differentiable(maybe_tensor):
     if not isinstance(maybe_tensor, torch.Tensor):
         return False
     return maybe_tensor.requires_grad
 
+
 def _any_differentiable(tensor_or_tuple_of_tensors):
     flat_args, _ = tree_unflatten(tensor_or_tuple_of_tensors)
     return any(tuple(map(_is_differentiable, flat_args)))
 
+
 def _wrap_tensor_for_grad(maybe_tensor, level):
     if not isinstance(maybe_tensor, torch.Tensor):
         return maybe_tensor
     return _wrap_for_grad(maybe_tensor, level)
 
+
 def _wrap_all_tensors(tensor_pytree, level):
     return tree_map(partial(_wrap_tensor_for_grad, level=level), tensor_pytree)
 
+
 def _as_tuple(val):
     if isinstance(val, tuple):
         return val
     return (val,)
 
 # Version of autograd.grad that handles outputs that don't depend on inputs
+
+
 def _autograd_grad(outputs, inputs, grad_outputs=None, retain_graph=False, create_graph=True):
     if grad_outputs is None:
         diff_outputs = tuple(out for out in outputs if out.requires_grad)
@@ -272,10 +280,12 @@ def wrapper(cotangents, retain_graph=True, create_graph=None):
     else:
         return results, wrapper
 
+
 def _safe_zero_index(x):
     assert len(x) == 1
     return x[0]
 
+
 def jacrev(f: Callable, argnums: Union[int, Tuple[int]] = 0, *, has_aux=False):
     """
     Computes the Jacobian of :attr:`f` with respect to the arg(s) at index
@@ -388,13 +398,13 @@ def wrapper_fn(*args):
         flat_output, output_spec = tree_flatten(output)
         if len(flat_output) == 0:
             raise RuntimeError(
-                f'jacrev(f, ...)(*args): expected f to return at least one Tensor, '
-                f'got no Tensors.')
+                'jacrev(f, ...)(*args): expected f to return at least one Tensor, '
+                'got no Tensors.')
         for out in flat_output:
             if isinstance(out, torch.Tensor):
                 continue
             raise RuntimeError(
-                f'jacrev(f, ...)(*args): expected f to only return Tensors as '
+                'jacrev(f, ...)(*args): expected f to only return Tensors as '
                 f'outputs, got {type(out)}')
         # NB: vjp already checks that all outputs are tensors
         # Step 1: Construct grad_outputs by splitting the standard basis
@@ -478,6 +488,7 @@ def wrapper_fn(*args):
 # - one of shape [3, 3] for the first output
 # - one of shape [   3] for the second output
 
+
 def _construct_standard_basis_for(tensors, tensor_numels):
     # This function:
     # - constructs a N=sum(tensor_numels) standard basis. i.e. an NxN identity matrix.
@@ -508,6 +519,7 @@ def _construct_standard_basis_for(tensors, tensor_numels):
                    for chunk, tensor in zip(chunks, tensors))
     return chunks
 
+
 def _validate_and_wrap_argnum(argnum, num_args):
     if not isinstance(argnum, int):
         raise RuntimeError(f'argnum must be int, got: {type(argnum)}')
@@ -517,6 +529,7 @@ def _validate_and_wrap_argnum(argnum, num_args):
         return argnum + num_args
     raise RuntimeError(f'Got argnum={argnum}, but only {num_args} positional inputs')
 
+
 def _check_unique_non_empty(argnums):
     if isinstance(argnums, tuple):
         if len(argnums) == 0:
@@ -524,6 +537,7 @@ def _check_unique_non_empty(argnums):
         if len(set(argnums)) != len(argnums):
             raise RuntimeError(f"argnums elements must be unique, got {argnums}")
 
+
 def _replace_args(old_args, new_args, argnums):
     if isinstance(argnums, int):
         if len(new_args) != 1:
@@ -534,10 +548,14 @@ def _replace_args(old_args, new_args, argnums):
             raise RuntimeError(
                 "new_args should have the same size as argnums. "
                 f"Argnums size {len(argnums)}, new_args size {len(new_args)}")
-        get_right_elem = lambda i : new_args[argnums.index(i)] if i in argnums else old_args[i]
+
+        def get_right_elem(i):
+            return new_args[argnums.index(i)] if i in argnums else old_args[i]
+
         return tuple(get_right_elem(i) for i in range(len(old_args)))
     raise RuntimeError(f'argnums must be int or Tuple[int, ...], got: {type(argnums)}')
 
+
 def _validate_and_wrap_argnums(argnums, num_args):
     if isinstance(argnums, int):
         return _validate_and_wrap_argnum(argnums, num_args)
@@ -545,6 +563,7 @@ def _validate_and_wrap_argnums(argnums, num_args):
         return tuple(_validate_and_wrap_argnum(argnum, num_args) for argnum in argnums)
     raise AssertionError("Should never get here")
 
+
 def _slice_argnums(args, argnums, as_tuple=True):
     if not isinstance(argnums, int) and not isinstance(argnums, tuple):
         raise RuntimeError(f'argnums must be int or Tuple[int, ...], got: {type(argnums)}')
@@ -557,6 +576,7 @@ def _slice_argnums(args, argnums, as_tuple=True):
             return args[argnums]
     return tuple(args[i] for i in argnums)
 
+
 def _argnums_partial(f, args, argnums):
     def f_wrapper(*wrapper_args):
         replaced_args = _replace_args(args, wrapper_args, argnums)
@@ -565,12 +585,15 @@ def f_wrapper(*wrapper_args):
     wrapper_args = wrapper_args if isinstance(wrapper_args, tuple) else (wrapper_args, )
     return (f_wrapper, wrapper_args)
 
+
 JVP_NESTING = 0
 
+
 @contextlib.contextmanager
 def noop():
     yield
 
+
 def assert_flat_tuple_of_tensors(elts, api, argname):
     if not isinstance(elts, tuple):
         raise RuntimeError(
@@ -585,6 +608,7 @@ def assert_flat_tuple_of_tensors(elts, api, argname):
         raise RuntimeError(
             f'{api}: Expected {argname} to be a non-empty tuple of Tensors.')
 
+
 def assert_output_is_tensor_or_tensors(output, api):
     if isinstance(output, torch.Tensor):
         return
@@ -602,6 +626,7 @@ def assert_output_is_tensor_or_tensors(output, api):
             f'{api}: Expected output of f to be a Tensor or Tensors, got '
             f'{type(out)} as an output')
 
+
 def assert_non_empty_list_of_tensors(output, api, argname):
     if len(output) == 0:
         raise RuntimeError(
@@ -613,19 +638,22 @@ def assert_non_empty_list_of_tensors(output, api, argname):
             f'{api}: Expected {argname} to only contain Tensors, got '
             f'{type(out)}')
 
+
 jvp_str = 'jvp(f, primals, tangents)'
 
+
 def safe_unpack_dual(dual, strict):
     primal, tangent = fwAD.unpack_dual(dual)
     if tangent is None:
         if strict:
             raise RuntimeError(
-                    f'jvp(f, primals, tangents, strict=True): '
-                    f'The output of f is independent of '
-                    f'the inputs. This is not allowed with strict=True.')
+                'jvp(f, primals, tangents, strict=True): '
+                'The output of f is independent of '
+                'the inputs. This is not allowed with strict=True.')
         tangent = torch.zeros_like(primal)
     return primal, tangent
 
+
 def jvp(f, primals, tangents, *, strict=False):
     if not isinstance(primals, tuple):
         raise RuntimeError(
@@ -670,12 +698,14 @@ def jvp(f, primals, tangents, *, strict=False):
         _grad_decrement_nesting()
         JVP_NESTING -= 1
 
+
 def safe_unflatten(tensor, dim, shape):
     if len(shape) == 0:
         assert tensor.shape[dim] == 1
         return tensor.squeeze(dim)
     return tensor.unflatten(dim, shape)
 
+
 def jacfwd(f, argnums=0):
     def wrapper_fn(*args):
         f_wrapper, primals = _argnums_partial(f, args, argnums)
@@ -707,9 +737,11 @@ def push_jvp(basis):
         return tree_unflatten(jac_outs_ins, spec)
     return wrapper_fn
 
+
 def hessian(f, argnums=0):
     return jacfwd(jacrev(f, argnums), argnums)
 
+
 def grad_and_value(func: Callable, argnums: argnums_t = 0, has_aux: bool = False) -> Callable:
     """
     Returns a function to compute a tuple of the gradient and primal, or
@@ -783,6 +815,7 @@ def wrapper(*args, **kwargs):
         return grad_input, output
     return wrapper
 
+
 def grad(func: Callable, argnums: argnums_t = 0, has_aux: bool = False) -> Callable:
     """``grad`` operator helps computing gradients of :attr:`func` with respect to the
     input(s) specified by :attr:`argnums`. This operator can be nested to
diff --git a/functorch/_src/fx_minifier.py b/functorch/_src/fx_minifier.py
index b1564bb16..e2195f266 100644
--- a/functorch/_src/fx_minifier.py
+++ b/functorch/_src/fx_minifier.py
@@ -1,35 +1,40 @@
+import subprocess
 import torch.fx as fx
 import copy
 import torch
 import math
 
+
 class ConcreteProp(torch.fx.Interpreter):
-        def run_node(self, n):
-                result = super().run_node(n)
+    def run_node(self, n):
+        result = super().run_node(n)
 
-                found_tensor = False
+        found_tensor = False
 
-                def extract_tensor_meta(obj):
-                        if isinstance(obj, torch.Tensor):
-                                nonlocal found_tensor
-                                found_tensor = True
-                                return obj
-                        else:
-                                return obj
+        def extract_tensor_meta(obj):
+            if isinstance(obj, torch.Tensor):
+                nonlocal found_tensor
+                found_tensor = True
+                return obj
+            else:
+                return obj
 
-                from torch.fx.node import map_aggregate
-                concrete_value = map_aggregate(result, extract_tensor_meta)
-                if found_tensor:
-                        n.meta['concrete_value'] = concrete_value
-                return result
+        from torch.fx.node import map_aggregate
+        concrete_value = map_aggregate(result, extract_tensor_meta)
+        if found_tensor:
+            n.meta['concrete_value'] = concrete_value
+        return result
+
+    def propagate(self, *args):
+        return super().run(*args)
 
-        def propagate(self, *args):
-                return super().run(*args)
 
 def _get_placeholders(graph):
     return list(filter(lambda x: x.op == 'placeholder', graph.nodes))
 
 # inplace modifies node/inps
+
+
 def _convert_node_to_placeholder(node, inps):
     node.op = 'placeholder'
     node.args = ()
@@ -42,13 +47,15 @@ def _convert_node_to_placeholder(node, inps):
         for tuple_user in list(node.users):
             _convert_node_to_placeholder(tuple_user, inps)
 
+
 def minimizer(fail_f: fx.GraphModule, inps, module_fails):
     """
     Minimizes a FX graph with given inputs, such that the resulting FX graph still returns True for module_fails.
 
     Does 2 main strategies:
     1. Truncates suffix: Removes some suffix from the graph and sets a new output.
-    2. Delta Debugging: Tries replacing half of the graph with inputs. If fails, tries replacing quarter of the graph, etc.
+    2. Delta Debugging: Tries replacing half of the graph with inputs. If fails,
+        tries replacing quarter of the graph, etc.
 
     >>> failing_function = fx.symbolic_trace(f)
     >>> minimize(failing_function, [torch.randn(5)], lambda fx_g, inps: fx_g(*inps))
@@ -94,7 +101,6 @@ def remove_suffix(cur_graph, cur_inps):
         print("FAIL: Could not remove suffix")
         return (cur_graph, cur_inps), False
 
-
     def remove_unused_inputs(cur_graph, cur_inps):
         assert graph_fails(cur_graph, cur_inps)
         ph_nodes = _get_placeholders(cur_graph)
@@ -139,7 +145,7 @@ def consolidate_placeholders(cur_graph):
                 new_node = new_graph.node_copy(node, lambda x: env[x])
                 env[node] = new_node
         return new_graph
-            
+
     def delta_debugging(cur_graph: fx.Graph, cur_inps):
         print("Strategy: Delta Debugging")
         assert graph_fails(cur_graph, cur_inps)
@@ -158,30 +164,36 @@ def delta_debugging(cur_graph: fx.Graph, cur_inps):
                         is_removing = True
                         _convert_node_to_placeholder(new_node, new_inps)
                 if not is_removing:
-                    continue 
+                    continue
                 new_graph = consolidate_placeholders(new_graph)
                 if graph_fails(new_graph, new_inps):
-                    print(f"SUCCESS: Removed ({start_range}:{end_range}] - Went from {starting_placeholders} placeholders to {len(_get_placeholders(new_graph))}")
+                    print(
+                        f"SUCCESS: Removed ({start_range}:{end_range}] - Went from {starting_placeholders} "
+                        "placeholders to {len(_get_placeholders(new_graph))}"
+                    )
                     return (new_graph, new_inps), True
             gap //= 2
 
         print("FAIL: Could not remove prefix")
         return (cur_graph, inps), False
 
-
-    print(f"###################")
+    print("###################")
     print(f"Current size: {len(failing_graph.nodes)}")
-    print(f"###################")
+    print("###################")
     while True:
         any_succeeded = False
-        for strategy in [remove_suffix, eliminate_dead_code, remove_unused_inputs, delta_debugging, eliminate_dead_code, remove_unused_inputs]:
+        strategies = [
+            remove_suffix, eliminate_dead_code, remove_unused_inputs,
+            delta_debugging, eliminate_dead_code, remove_unused_inputs
+        ]
+        for strategy in strategies:
             out = strategy(copy.deepcopy(failing_graph), inps[:])
             (cur_graph, cur_inps), succeeded = out
             if succeeded:
                 print()
-                print(f"###################")
+                print("###################")
                 print(f"Current size: {len(cur_graph.nodes)}")
-                print(f"###################")
+                print("###################")
                 failing_graph = cur_graph
                 inps = cur_inps
                 any_succeeded = True
@@ -193,8 +205,9 @@ def delta_debugging(cur_graph: fx.Graph, cur_inps):
     print([i.shape for i in inps])
     return failing_fx, inps
 
+
 def check_nvfuser_subprocess(f, inps):
-    import subprocess
+
     f.to_folder("temp")
     with open("_temp.py", 'w') as fil:
         fil.write(f'''
diff --git a/functorch/_src/make_functional.py b/functorch/_src/make_functional.py
index a62036453..6075466aa 100644
--- a/functorch/_src/make_functional.py
+++ b/functorch/_src/make_functional.py
@@ -9,11 +9,12 @@
 from torch import Tensor
 from typing import List, Tuple
 import copy
-import warnings
 
 # Utilities to make nn.Module "functional"
 # In particular the goal is to be able to provide a function that takes as input
 # the parameters and evaluate the nn.Module using fixed inputs.
+
+
 def _del_nested_attr(obj: nn.Module, names: List[str]) -> None:
     """
     Deletes the attribute specified by the given list of names.
@@ -25,6 +26,7 @@ def _del_nested_attr(obj: nn.Module, names: List[str]) -> None:
     else:
         _del_nested_attr(getattr(obj, names[0]), names[1:])
 
+
 def _set_nested_attr(obj: nn.Module, names: List[str], value: Tensor) -> None:
     """
     Set the attribute specified by the given list of names to value.
@@ -36,12 +38,14 @@ def _set_nested_attr(obj: nn.Module, names: List[str], value: Tensor) -> None:
     else:
         _set_nested_attr(getattr(obj, names[0]), names[1:], value)
 
+
 def _get_nested_attr(obj: nn.Module, names: List[str]) -> None:
     if len(names) == 1:
         return getattr(obj, names[0])
     else:
         _get_nested_attr(getattr(obj, names[0]), names[1:])
 
+
 def extract_weights(mod: nn.Module) -> Tuple[Tuple[Tensor, ...], List[str]]:
     """
     This function removes all the Parameters from the model and
@@ -63,6 +67,7 @@ def extract_weights(mod: nn.Module) -> Tuple[Tuple[Tensor, ...], List[str]]:
     params = tuple(p for p in orig_params)
     return params, names
 
+
 def load_weights(mod: nn.Module, names: List[str], params: Tuple[Tensor, ...], as_params=False) -> None:
     """
     Reload a set of weights so that `mod` can be used again to perform a forward pass.
@@ -75,6 +80,7 @@ def load_weights(mod: nn.Module, names: List[str], params: Tuple[Tensor, ...], a
         _del_nested_attr(mod, name.split("."))
         _set_nested_attr(mod, name.split("."), p)
 
+
 def _swap_state(mod: nn.Module, split_names: List[str], elems):
     result = []
     for split_name, elem in zip(split_names, elems):
@@ -83,6 +89,7 @@ def _swap_state(mod: nn.Module, split_names: List[str], elems):
         _set_nested_attr(mod, split_name, elem)
     return result
 
+
 def extract_buffers(mod: nn.Module) -> Tuple[Tuple[Tensor, ...], List[str]]:
     orig_params = tuple(mod.buffers())
     # Remove all the parameters in the model
@@ -195,13 +202,16 @@ def fun(weights, buffers, data):
 
     return weights, buffers, fun, weight_descriptors, buf_descriptors
 
+
 def make_split_names(lst):
     return [name.split('.') for name in lst]
 
+
 class FunctionalModuleWithBuffers(nn.Module):
     """
     This is the callable object returned by :func:`make_functional_with_buffers`.
     """
+
     def __init__(self, stateless_model, param_names, buffer_names):
         super(FunctionalModuleWithBuffers, self).__init__()
         self.stateless_model = stateless_model
@@ -233,10 +243,12 @@ def forward(self, params, buffers, *args, **kwargs):
             # Remove the loaded state on self.stateless_model
             _swap_state(self.stateless_model, self.split_names, old_state)
 
+
 class FunctionalModule(nn.Module):
     """
     This is the callable object returned by :func:`make_functional`.
     """
+
     def __init__(self, stateless_model, param_names):
         super(FunctionalModule, self).__init__()
         self.stateless_model = stateless_model
@@ -259,6 +271,7 @@ def forward(self, params, *args, **kwargs):
             # Remove the loaded state on self.stateless_model
             _swap_state(self.stateless_model, self.split_names, old_state)
 
+
 def make_functional(model: nn.Module):
     """make_functional(model) -> func, params
 
diff --git a/functorch/_src/memory_efficient_op_authoring.py b/functorch/_src/memory_efficient_op_authoring.py
index 3ec998a2b..cd9c62c61 100644
--- a/functorch/_src/memory_efficient_op_authoring.py
+++ b/functorch/_src/memory_efficient_op_authoring.py
@@ -1,7 +1,4 @@
-from functorch._src.python_key import pythonkey_decompose
 import torch
-from torch import fx
-from torch.fx.proxy import GraphAppendingTracer
 from typing import Iterable
 from .aot_autograd import compiled_function, partition_with_recompute_fwd_in_bwd
 
diff --git a/functorch/_src/nnc_compile.py b/functorch/_src/nnc_compile.py
index 41de5d400..eaca7ea3e 100644
--- a/functorch/_src/nnc_compile.py
+++ b/functorch/_src/nnc_compile.py
@@ -6,8 +6,6 @@
 
 import time
 import torch
-import torch.nn as nn
-import torch.nn.functional as F
 import torch._C._te as te
 import torch.fx as fx
 import torch.utils._pytree as pytree
@@ -16,9 +14,10 @@
 import operator
 import functools
 
+
 def truncate(model, k):
     model = fx.symbolic_trace(model)
-    new_graph= fx.Graph()
+    new_graph = fx.Graph()
     env = {}
 
     cnt = 0
@@ -33,6 +32,8 @@ def truncate(model, k):
     return fx.GraphModule(model, new_graph)
 
 # NNC Lowering Pass
+
+
 def remove_args(model: torch.nn.Module, args):
     fx_model = fx.symbolic_trace(model)
     for node in fx_model.graph.nodes:
@@ -42,15 +43,18 @@ def remove_args(model: torch.nn.Module, args):
     fx_model.recompile()
     return fx_model
 
+
 def get_dim_args(dims):
     dim_args = []
     for dim in dims:
         dim_args.append(te.DimArg(te.ExprHandle.int(dim), 'i' + str(len(dim_args))))
     return dim_args
 
+
 def get_te_shapes(shape):
     return [te.ExprHandle.int(i) for i in shape]
 
+
 def to_expr(x):
     if isinstance(x, int):
         return te.ExprHandle.int(x)
@@ -59,6 +63,7 @@ def to_expr(x):
     else:
         raise RuntimeError(f"type {type(x)} not supported")
 
+
 def get_nnc_type(dtype):
     if dtype == torch.float:
         return te.Dtype.Float
@@ -74,29 +79,35 @@ def get_nnc_type(dtype):
         raise RuntimeError(f"type nyi {dtype}")
 
 
-lowering_functions = { }
+lowering_functions = {}
+
+
 def index_or_broadcast(shape, *args):
     out = []
     for idx, arg in enumerate(args):
-        if idx >= len(shape): continue
+        if idx >= len(shape):
+            continue
         if shape[idx] == 1:
             out.append(to_expr(0))
         else:
             out.append(arg)
     return out
 
+
 def ones_like_lower(name, out_shape, inp_shapes, args):
     def f(*idxs):
         return to_expr(1.0)
     res = te.Compute(name, get_dim_args(out_shape), f)
     return res
 
+
 def zeros_like_lower(name, out_shape, inp_shapes, args):
     def f(*idxs):
         return to_expr(0.0)
     res = te.Compute(name, get_dim_args(out_shape), f)
     return res
 
+
 def full_like_lower(name, out_shape, inp_shapes, args):
     def f(*idxs):
         return to_expr(args[1])
@@ -106,9 +117,11 @@ def f(*idxs):
 
 def prod(x, start=1):
     t = start
-    for i in x: t *= i
+    for i in x:
+        t *= i
     return t
 
+
 def encode_idxs(shape, idxs):
     assert(len(shape) == len(idxs))
     cur = 1
@@ -118,10 +131,12 @@ def encode_idxs(shape, idxs):
         cur *= dim
     return out
 
+
 def reshape_lower(name, out_shape, inp_shapes, args):
     X, shape = args
     start_shape = list(inp_shapes[0][0])
     end_shape = out_shape
+
     def get_orig_idxs(idxs):
         absolute_new = encode_idxs(end_shape, idxs)
         new_idxs = []
@@ -151,11 +166,13 @@ def f(*idxs):
 #     res = te.Compute(name, get_dim_args(out_shape), f)
 #     return res
 
+
 def dot_lower(name, out_shape, inp_shapes, args):
     mul_te = te.lower('aten::mul', list(args), get_te_shapes(inp_shapes[0][0]), get_nnc_type(inp_shapes[0][1]))
     res = te.lower('aten::sum', [mul_te.buf()], get_te_shapes(out_shape), get_nnc_type(inp_shapes[0][1]))
     return (res.buf(), [mul_te.stmt(), res.stmt()])
 
+
 def mv_lower(name, out_shape, inp_shapes, args):
     A = args[0]
     B = args[1]
@@ -170,35 +187,42 @@ def mv_lower(name, out_shape, inp_shapes, args):
     s = torch._C._te.ExternalCall(C, "nnc_aten_mv", [A, B], [])
     return C, [s]
 
+
 def digamma_lower(name, out_shape, inp_shapes, args):
     out = te.BufHandle('out', get_te_shapes(out_shape), get_nnc_type(inp_shapes[0][1]))
     s = te.ExternalCall(out, "nnc_aten_digamma", [args[0]], [])
     return out, [s]
 
+
 def ger_lower(name, out_shape, inp_shapes, args):
-    A = args[0]
-    B = args[1]
     A_len = inp_shapes[0][0][0]
     B_len = inp_shapes[1][0][0]
     A_squeeze = te.lower('aten::unsqueeze', [args[0], 1], get_te_shapes([A_len, 1]), get_nnc_type(inp_shapes[0][1]))
     B_squeeze = te.lower('aten::unsqueeze', [args[1], 0], get_te_shapes([1, B_len]), get_nnc_type(inp_shapes[1][1]))
-    out = te.lower('aten::mul', [A_squeeze.buf(), B_squeeze.buf()], get_te_shapes(out_shape), get_nnc_type(inp_shapes[0][1]))
+    out = te.lower('aten::mul', [A_squeeze.buf(), B_squeeze.buf()],
+                   get_te_shapes(out_shape), get_nnc_type(inp_shapes[0][1]))
     return out.buf(), [A_squeeze.stmt(), B_squeeze.stmt(), out.stmt()]
 
+
 def triangular_solve_lower(name, out_shape, inp_shapes, args):
     A = args[0]
     B = args[1]
     C = torch._C._te.BufHandle('C', get_te_shapes(out_shape[0]), get_nnc_type(inp_shapes[0][1]))
-    s = torch._C._te.ExternalCall(C, "nnc_aten_triangular_solve", [A, B], [to_expr(args[2]), to_expr(args[3]), to_expr(args[4])])
+    s = torch._C._te.ExternalCall(C, "nnc_aten_triangular_solve", [A, B], [
+                                  to_expr(args[2]), to_expr(args[3]), to_expr(args[4])])
     return (C, None), [s]
 
+
 def binary_cross_entropy_lower(name, out_shape, inp_shapes, args):
     self_ = args[0]
     target = args[1]
-    if args[2] != None or args[3] == 2:
+    if args[2] is not None or args[3] == 2:
         raise RuntimeError(f"weight={args[2]} and reduction={args[3]} not supported")
+
     def f(*idxs):
-        return to_expr(0.0) - (self_.load(idxs).log() * target.load(idxs) + (to_expr(1.0) - target.load(idxs)) * (to_expr(1.0) - self_.load(idxs)).log())
+        out = (to_expr(1.0) - target.load(idxs)) * (to_expr(1.0) - self_.load(idxs)).log()
+        return to_expr(0.0) - (self_.load(idxs).log() * target.load(idxs) + out)
+
     val = te.Compute(name, get_dim_args(inp_shapes[0][0]), f)
     if args[3] == 0:
         return val.buf(), [val.stmt()]
@@ -208,15 +232,18 @@ def f(*idxs):
 
 def binary_cross_entropy_with_logits_lower(name, out_shape, inp_shapes, args):
     pred = te.lower('aten::sigmoid', [args[0]], get_te_shapes(inp_shapes[0][0]), get_nnc_type(inp_shapes[0][1]))
-    loss_buf, loss_stmts = binary_cross_entropy_lower('binary_cross_entropy', out_shape, list(inp_shapes) + [None], [pred, args[1], args[3], args[4]])
+    loss_buf, loss_stmts = binary_cross_entropy_lower(
+        'binary_cross_entropy', out_shape, list(inp_shapes) + [None], [pred, args[1], args[3], args[4]])
     return loss_buf, [pred.stmt()] + loss_stmts
 
+
 def detach_lower(name, out_shape, inp_shapes, args):
     return args[0], []
 
 # def clone_lower(name, out_shape, inp_shapes, args):
 #     return args[0], []
 
+
 lowering_functions[torch.ops.aten.full_like] = full_like_lower
 lowering_functions[torch.ops.aten.zeros_like] = zeros_like_lower
 lowering_functions[torch.ops.aten.ones_like] = ones_like_lower
@@ -234,7 +261,6 @@ def detach_lower(name, out_shape, inp_shapes, args):
 # lowering_functions[torch.ops.aten.clone] = clone_lower
 
 
-
 func_to_aten = {
     operator.add: torch.ops.aten.add,
     operator.mul: torch.ops.aten.mul,
@@ -249,6 +275,7 @@ def process_shape(x):
         return torch.Size([1])
     return x
 
+
 def map_node_meta(f, node_meta):
     if isinstance(node_meta, TensorMetadata):
         return f(node_meta)
@@ -258,8 +285,15 @@ def map_node_meta(f, node_meta):
         return list([map_node_meta(f, i) for i in node_meta])
     return f(node_meta)
 
+
 def lower_function(node, op, nnc_args, args):
-    inp_shapes = fx.node.map_aggregate(args, lambda arg: (process_shape(arg.meta['tensor_meta'].shape), arg.meta['tensor_meta'].dtype) if isinstance(arg, fx.Node) and 'tensor_meta' in arg.meta else None)
+    inp_shapes = fx.node.map_aggregate(
+        args,
+        lambda arg: (
+            process_shape(arg.meta['tensor_meta'].shape),
+            arg.meta['tensor_meta'].dtype
+        ) if isinstance(arg, fx.Node) and 'tensor_meta' in arg.meta else None
+    )
     out_shape = map_node_meta(lambda x: process_shape(x.shape), node.meta['tensor_meta'])
     if op in lowering_functions:
         out = lowering_functions[op](node.name, out_shape, inp_shapes, nnc_args)
@@ -274,6 +308,8 @@ def lower_function(node, op, nnc_args, args):
         return out[0], out[1]
 
 # This will not work properly in the presence of aliasing/views
+
+
 def remove_inplace(fx_model: fx.GraphModule) -> torch.nn.Module:
     new_map = {}
     for node in fx_model.graph.nodes:
@@ -283,7 +319,8 @@ def remove_inplace(fx_model: fx.GraphModule) -> torch.nn.Module:
             new_map[node.args[0]] = node
     return fx_model
 
-def nnc_compile(fx_model: fx.GraphModule, example_inputs, get_loopnest = False) -> torch.nn.Module:
+
+def nnc_compile(fx_model: fx.GraphModule, example_inputs, get_loopnest=False) -> torch.nn.Module:
     """
     nnc_compile(model, example_inputs) returns a function with the same args
     as `model.forward`, with an extra argument corresponding to where the
@@ -291,7 +328,6 @@ def nnc_compile(fx_model: fx.GraphModule, example_inputs, get_loopnest = False)
     tensors with the same shapes as example_inputs), and passes them to an
     NNC executor.
     """
-    t = fx_model.graph.flatten_inps(*example_inputs)
     ShapeProp(fx_model).propagate(*fx_model.graph.flatten_inps(*example_inputs))
     fx_model = remove_inplace(fx_model)
 
@@ -299,18 +335,14 @@ def nnc_compile(fx_model: fx.GraphModule, example_inputs, get_loopnest = False)
     # of an NNC computation.
     env = {}
 
-
     def get_te_type(node):
         return get_nnc_type(node.meta['tensor_meta'].dtype)
 
-    def gen_compute(args):
-        te_args = [env[arg.name] for arg in args]
-
-    def lookup_env(l):
-        res = fx.node.map_aggregate(l, lambda x: env[x.name] if isinstance(x, fx.Node) else x)
+    def lookup_env(inpt):
+        res = fx.node.map_aggregate(inpt, lambda x: env[x.name] if isinstance(x, fx.Node) else x)
         return res
 
-    def fetch_attr(target : str):
+    def fetch_attr(target: str):
         target_atoms = target.split('.')
         attr_itr = fx_model
         for i, atom in enumerate(target_atoms):
@@ -357,7 +389,8 @@ def fetch_attr(target : str):
                 continue
         elif node.op == 'output':
             args = node.args
-            args = pytree.tree_map(lambda x: list(x) if isinstance(x, fx.immutable_collections.immutable_list) else x, args)
+            args = pytree.tree_map(lambda x: list(x) if isinstance(
+                x, fx.immutable_collections.immutable_list) else x, args)
             flat_args, _ = pytree.tree_flatten(list(args))
             te_args = lookup_env(flat_args)
             outs = (list(te_args), [(i.meta['tensor_meta'].shape, i.meta['tensor_meta'].dtype) for i in flat_args])
@@ -366,7 +399,7 @@ def fetch_attr(target : str):
             # attributes and pass them in as inputs to NNC.
             module_attrs.append(node)
             shapes = get_te_shapes(process_shape(node.meta['tensor_meta'].shape))
-            placeholder = te.BufHandle(node.name, shapes,  get_te_type(node))
+            placeholder = te.BufHandle(node.name, shapes, get_te_type(node))
             env[node.name] = placeholder
             attr_bufs.append(placeholder)
             inputs_or_attrs.add(placeholder)
@@ -374,7 +407,6 @@ def fetch_attr(target : str):
             print(node.op, node.target)
             raise RuntimeError("not yet implemented")
 
-
     if len(compute_stmts) == 0:
         raise RuntimeError("Doesn't support compiling empty")
 
@@ -388,7 +420,8 @@ def fetch_attr(target : str):
     loopnest.simplify()
     loopnest.prepare_for_codegen()
     stmt = te.simplify(loopnest.root_stmt())
-    cg = te.construct_codegen('llvm', stmt, [te.BufferArg(x) for x in [env[i.name] for i in module_attrs] + inputs + buf_outs])
+    cg = te.construct_codegen('llvm', stmt, [te.BufferArg(x)
+                              for x in [env[i.name] for i in module_attrs] + inputs + buf_outs])
 
     module_stuff = [fetch_attr(i.target).contiguous().data for i in module_attrs]
 
@@ -400,7 +433,10 @@ def fetch_attr(target : str):
             ph_to_inp_map[outs[idx][0]] = attr_bufs.index(outs[idx][0])
 
     def get_outs(inps):
-        return [inps[ph_to_inp_map[buf]] if buf in inputs_or_attrs else torch.empty(shape, dtype=dtype) for buf, (shape,dtype) in outs]
+        return [
+            inps[ph_to_inp_map[buf]] if buf in inputs_or_attrs else torch.empty(shape, dtype=dtype)
+            for buf, (shape, dtype) in outs
+        ]
 
     def f(*inps):
         inps = fx_model.graph.flatten_inps(*inps)
@@ -415,6 +451,7 @@ def f(*inps):
         return results
     return f
 
+
 def make_nnc(f):
     @functools.wraps(f)
     def wrapped(*args):
@@ -425,6 +462,7 @@ def wrapped(*args):
 
     return wrapped
 
+
 def get_ops(fx_model: fx.GraphModule):
     vals = set()
     for node in fx_model.graph.nodes:
@@ -433,7 +471,6 @@ def get_ops(fx_model: fx.GraphModule):
     return vals
 
 
-
 ################################
 # Example usage and Benchmarking
 ################################
@@ -444,11 +481,12 @@ def bench(f, warmup=3, iters=1000):
     begin = time.time()
     for _ in range(iters):
         f()
-    print(time.time()-begin)
+    print(time.time() - begin)
+
 
 if __name__ == '__main__':
     def f(a, b):
-        return (torch.cos(a)* torch.sin(b))[:2000]
+        return (torch.cos(a) * torch.sin(b))[:2000]
 
     mod = fx.symbolic_trace(f)
     inps = (torch.randn(5000), torch.randn(5000))
diff --git a/functorch/_src/operator_authoring.py b/functorch/_src/operator_authoring.py
index cfcf72cff..ef924cb28 100644
--- a/functorch/_src/operator_authoring.py
+++ b/functorch/_src/operator_authoring.py
@@ -13,7 +13,12 @@
 FOLD_ALIASES = True
 _SHAPE_TYPES = {"one", "other"}
 _STRIDE_TYPES = {"zero", "one", "contiguous", "transposed_contiguous", "as_arg"}
-_identity = lambda x: x
+
+
+def _identity(x):
+    return x
+
+
 _TORCH_TO_EXPR_MAP = {
     "sin": _te.sin,
     "cos": _te.cos,
diff --git a/functorch/_src/python_key.py b/functorch/_src/python_key.py
index 4f751c2a7..595627f29 100644
--- a/functorch/_src/python_key.py
+++ b/functorch/_src/python_key.py
@@ -5,15 +5,12 @@
 # LICENSE file in the root directory of this source tree.
 from dataclasses import dataclass
 import functools
-from typing import Any, Dict, NamedTuple, Optional, Set, Tuple, List, Callable, Union
+from typing import Any, Dict, Optional, Tuple, Callable, Union
 import torch
 from torch._C import _disabled_torch_function_impl
-from torch.fx.node import map_aggregate
 import torch.utils._pytree as pytree
 from torch.fx import Tracer, GraphModule
 import torch.fx as fx
-import torch.fx._pytree as fx_pytree
-from torch import Tensor
 from .nnc_compile import nnc_compile
 from .decompositions import decomposition_table
 from enum import Enum
@@ -24,6 +21,7 @@
 USE_DECOMPOSE = False
 USE_META = False
 
+
 @contextmanager
 def pythonkey_decompose():
     global USE_DECOMPOSE
@@ -43,6 +41,7 @@ def pythonkey_meta():
     finally:
         USE_META = False
 
+
 class PythonTensor(torch.Tensor):
     elem: torch.Tensor
 
@@ -72,6 +71,7 @@ def __repr__(self):
         return f"PythonTensor({self.elem})"
 
     __torch_function__ = _disabled_torch_function_impl
+
     @classmethod
     def __torch_dispatch__(cls, func, types, args=(), kwargs=None):
         if func in decomposition_table and USE_DECOMPOSE:
@@ -84,7 +84,8 @@ def unwrap_tensor(e):
             return e.elem if isinstance(e, PythonTensor) else e
 
         # Used to infer the output device
-        input_devices = list(set([i.device for i in pytree.tree_flatten(args)[0] + pytree.tree_flatten(kwargs)[0] if isinstance(i, PythonTensor)]))
+        input_devices = list(set([i.device for i in pytree.tree_flatten(args)[0] +
+                             pytree.tree_flatten(kwargs)[0] if isinstance(i, PythonTensor)]))
         assert len(input_devices) == 1
         output_device = input_devices[0]
         proxy_args = pytree.tree_map(unwrap_proxy, args)
@@ -94,17 +95,22 @@ def unwrap_tensor(e):
         kwargs = pytree.tree_map(unwrap_tensor, kwargs)
         try:
             real_out = func(*args, **kwargs)
-        except NotImplementedError as e:
-            args = pytree.tree_map(lambda x: torch.ones_like(x, device=output_device) if isinstance(x, torch.Tensor) else x, args)
-            kwargs = pytree.tree_map(lambda x: torch.ones_like(x, device=output_device) if isinstance(x, torch.Tensor) else x, kwargs)
+        except NotImplementedError:
+            args = pytree.tree_map(lambda x: torch.ones_like(x, device=output_device)
+                                   if isinstance(x, torch.Tensor) else x, args)
+            kwargs = pytree.tree_map(lambda x: torch.ones_like(x, device=output_device)
+                                     if isinstance(x, torch.Tensor) else x, kwargs)
             real_out = func(*args, **kwargs)
 
         def wrap_with_proxy(e, proxy):
-            # Some ops (like native_batch_norm_backward) return undefined tensors that get converted into None in python.
-            # As the function signature expects tensors, if we directly return these None tensors back to C++, we'll error.
+            # Some ops (like native_batch_norm_backward) return undefined tensors that get
+            # converted into None in python.
+            # As the function signature expects tensors, if we directly return these None
+            # tensors back to C++, we'll error.
             if e is None:
                 e = torch.empty(())
-            # Currently assuming that all inputs to an op are the same device - not totally sure that's true
+            # Currently assuming that all inputs to an op are the same device - not totally
+            # sure that's true
             if type(e) == torch.Tensor:
                 return PythonTensor(e, proxy, output_device)
             else:
@@ -118,12 +124,14 @@ def wrap_with_proxy(e, proxy):
         else:
             return real_out
 
+
 class PythonKeyTracer(Tracer):
     def __init__(self):
         super().__init__()
 
-
-    def call_module(self, m: torch.nn.Module, forward: Callable[..., Any], args : Tuple[Any, ...], kwargs : Dict[str, Any]) -> Any:
+    def call_module(
+        self, m: torch.nn.Module, forward: Callable[..., Any], args: Tuple[Any, ...], kwargs: Dict[str, Any]
+    ) -> Any:
         return forward(*args, **kwargs)
 
     def _module_getattr(self, attr, attr_val, parameter_proxy_cache):
@@ -145,7 +153,7 @@ def create_arg(self, a: Any):
             for n, p in self.root.named_parameters():
                 if a is p:
                     return self.create_node('get_attr', n, (), {})
-            qualname : Optional[str] = None
+            qualname: Optional[str] = None
 
             if not qualname:
                 i = 0
@@ -160,14 +168,18 @@ def create_arg(self, a: Any):
         return super().create_arg(a)
 
 
-def pythonkey_trace(root : Union[torch.nn.Module, Callable], concrete_args: Optional[Dict[str, Any]] = None) -> GraphModule:
+def pythonkey_trace(
+    root: Union[torch.nn.Module, Callable], concrete_args: Optional[Dict[str, Any]] = None
+) -> GraphModule:
     tracer = PythonKeyTracer()
     graph = tracer.trace(root, concrete_args)
     name = root.__class__.__name__ if isinstance(root, torch.nn.Module) else root.__name__
     return GraphModule(tracer.root, graph, name)
 
+
 def wrap_key(f, inps):
     flat_inps, inp_spec = pytree.tree_flatten(inps)
+
     @functools.wraps(f)
     def wrapped(*args):
         flat_args, args_spec = pytree.tree_flatten(args)
@@ -188,6 +200,7 @@ def wrapped(*args):
 
     return wrapped
 
+
 def make_fx(f):
     @functools.wraps(f)
     def wrapped(*args):
@@ -197,6 +210,7 @@ def wrapped(*args):
 
     return wrapped
 
+
 @dataclass(eq=True, frozen=True)
 class TensorSpec:
     shape: Tuple[int, ...]
@@ -204,15 +218,18 @@ class TensorSpec:
     dtype: torch.dtype
     device: torch.device
 
+
 @dataclass(eq=True, frozen=True)
 class ConcreteValueSpec:
     value: Any
 
+
 @dataclass(eq=True, frozen=True)
 class SpecializationKey:
     func: Callable
     specs: Tuple[Union[TensorSpec, ConcreteValueSpec], ...]
 
+
 def get_spec(arg):
     if isinstance(arg, torch.Tensor):
         return TensorSpec(
@@ -222,17 +239,21 @@ def get_spec(arg):
             arg.device)
     return ConcreteValueSpec(arg)
 
+
 def construct_specialization_key(f, args):
     flat_args, _ = pytree.tree_flatten(args)
     return SpecializationKey(f, tuple(get_spec(arg) for arg in flat_args))
 
+
 nnc_jit_cache: Dict[Callable, Dict[SpecializationKey, Callable]] = {}
 
+
 class RetrievalStatus(Enum):
     Success = 0
     UnknownFunc = 1
     UnknownSpecialization = 2
 
+
 def retrieve_from_cache(f, key):
     if f not in nnc_jit_cache:
         return RetrievalStatus.UnknownFunc, None
@@ -241,14 +262,17 @@ def retrieve_from_cache(f, key):
         return RetrievalStatus.UnknownSpecialization, None
     return RetrievalStatus.Success, cache_for_f[key]
 
+
 def add_to_cache(f, key, compiled_f):
     if f not in nnc_jit_cache:
         nnc_jit_cache[f] = {key: compiled_f}
     else:
         nnc_jit_cache[f][key] = compiled_f
 
-def nnc_jit(f, static_argnums=None, skip_specialization = False):
+
+def nnc_jit(f, static_argnums=None, skip_specialization=False):
     local_cache = None
+
     @functools.wraps(f)
     def compiled(*args):
         nonlocal local_cache, static_argnums
@@ -283,6 +307,7 @@ def compiled(*args):
         return compiled_f(*args)
     return compiled
 
+
 def make_nnc(f):
     @functools.wraps(f)
     def wrapped(*args):
diff --git a/functorch/_src/pytree_hacks.py b/functorch/_src/pytree_hacks.py
index db1cf6e33..a7a035042 100644
--- a/functorch/_src/pytree_hacks.py
+++ b/functorch/_src/pytree_hacks.py
@@ -13,6 +13,8 @@
 # I can't figure out how to do that, so we are turning all of them
 # into normal Tuples for now (this is what vmap used to do anyways).
 # We probably want some special behavior for named tuples?
+
+
 def tree_flatten_hack(pytree):
     if _pytree._is_leaf(pytree) and not isinstance(pytree, tuple):
         return [pytree], _pytree.LeafSpec()
@@ -35,17 +37,19 @@ def tree_flatten_hack(pytree):
 
     return result, _pytree.TreeSpec(typ, context, children_specs)
 
+
 def tree_map_(fn_, pytree):
     flat_args, _ = tree_flatten(pytree)
     [fn_(arg) for arg in flat_args]
     return pytree
 
+
 class PlaceHolder():
     def __repr__(self):
         return '*'
 
+
 def treespec_pprint(spec):
     leafs = [PlaceHolder() for _ in range(spec.num_leaves)]
     result = tree_unflatten(leafs, spec)
     return repr(result)
-
diff --git a/functorch/_src/top_operators_github_usage.py b/functorch/_src/top_operators_github_usage.py
index fddebce06..9161f98d6 100644
--- a/functorch/_src/top_operators_github_usage.py
+++ b/functorch/_src/top_operators_github_usage.py
@@ -598,6 +598,7 @@
     'view_as',
 ]
 
+
 def get_nn_functional_top_list():
     top_nn_functional_ = {k: v for k, v in top_nn_functional}
     for _, count, functional_name in top_nn_module:
@@ -614,6 +615,7 @@ def get_nn_functional_top_list():
     top_nn_functional_.sort(key=lambda x: x[1], reverse=True)
     return top_nn_functional_
 
+
 usage_count = {}
 for k, v in get_nn_functional_top_list():
     usage_count[k] = v
diff --git a/functorch/_src/vmap.py b/functorch/_src/vmap.py
index 849ae628d..696928dd8 100644
--- a/functorch/_src/vmap.py
+++ b/functorch/_src/vmap.py
@@ -11,7 +11,6 @@
 from torch.utils._pytree import tree_flatten, tree_unflatten, _broadcast_to_and_flatten, TreeSpec
 from .pytree_hacks import tree_flatten_hack, tree_map_
 from functools import partial
-import warnings
 
 from functorch._C import (
     _add_batch_dim,
@@ -24,6 +23,8 @@
 out_dims_t = Union[int, Tuple[int, ...]]
 
 # Checks that all args-to-be-batched have the same batch dim size
+
+
 def _validate_and_get_batch_size(
         flat_in_dims: List[Optional[int]],
         flat_args: List) -> int:
@@ -35,6 +36,7 @@ def _validate_and_get_batch_size(
             f'dimension, got sizes {batch_sizes} for the mapped dimension')
     return batch_sizes[0]
 
+
 def _num_outputs(batched_outputs: Union[Tensor, Tuple[Tensor, ...]]) -> int:
     if isinstance(batched_outputs, tuple):
         return len(batched_outputs)
@@ -42,6 +44,8 @@ def _num_outputs(batched_outputs: Union[Tensor, Tuple[Tensor, ...]]) -> int:
 
 # If value is a tuple, check it has length `num_elements`.
 # If value is not a tuple, make a tuple with `value` repeated `num_elements` times
+
+
 def _as_tuple(value: Any, num_elements: int, error_message_lambda: Callable[[], str]) -> Tuple:
     if not isinstance(value, tuple):
         return (value,) * num_elements
@@ -49,7 +53,10 @@ def _as_tuple(value: Any, num_elements: int, error_message_lambda: Callable[[],
         raise ValueError(error_message_lambda())
     return value
 
-def _process_batched_inputs(in_dims: in_dims_t, args: Tuple, func: Callable) -> Tuple[int, List[Any], List[Any], TreeSpec]:
+
+def _process_batched_inputs(
+    in_dims: in_dims_t, args: Tuple, func: Callable
+) -> Tuple[int, List[Any], List[Any], TreeSpec]:
     if not isinstance(in_dims, int) and not isinstance(in_dims, tuple):
         raise ValueError(
             f'vmap({_get_name(func)}, in_dims={in_dims}, ...)(<inputs>): '
@@ -95,6 +102,8 @@ def _process_batched_inputs(in_dims: in_dims_t, args: Tuple, func: Callable) ->
 
 # Creates BatchedTensors for every Tensor in arg that should be batched.
 # Returns the (potentially) batched arguments and the batch_size.
+
+
 def _create_batched_inputs(
         flat_in_dims: List[Any], flat_args: List[Any], vmap_level: int, args_spec) -> Tuple:
     # See NOTE [Ignored _remove_batch_dim, _add_batch_dim]
@@ -104,6 +113,8 @@ def _create_batched_inputs(
     return tree_unflatten(batched_inputs, args_spec)
 
 # Undos the batching (and any batch dimensions) associated with the `vmap_level`.
+
+
 def _unwrap_batched(
         batched_outputs: Union[Tensor, Tuple[Tensor, ...]],
         out_dims: out_dims_t,
@@ -144,6 +155,7 @@ def incompatible_error():
     ]
     return tree_unflatten(flat_outputs, output_spec)
 
+
 def _check_int(x, func, out_dims):
     if isinstance(x, int):
         return
@@ -152,11 +164,13 @@ def _check_int(x, func, out_dims):
         f'an int or a python collection of ints representing where in the outputs the '
         f'vmapped dimension should appear.')
 
+
 def _check_out_dims_is_int_or_int_pytree(out_dims: out_dims_t, func: Callable) -> None:
     if isinstance(out_dims, int):
         return
     tree_map_(partial(_check_int, func=func, out_dims=out_dims), out_dims)
 
+
 def _get_name(func: Callable):
     if hasattr(func, '__name__'):
         return func.__name__
@@ -169,6 +183,8 @@ def _get_name(func: Callable):
 # vmap(func)(inputs) wraps all Tensor inputs to be batched in BatchedTensors,
 # sends those into func, and then unwraps the output BatchedTensors. Operations
 # on BatchedTensors perform the batched operations that the user is asking for.
+
+
 def vmap(func: Callable, in_dims: in_dims_t = 0, out_dims: out_dims_t = 0) -> Callable:
     """
     vmap is the vectorizing map; ``vmap(func)`` returns a new function that
diff --git a/op_analysis/gen_data.py b/op_analysis/gen_data.py
index 2c846874a..c455d5b85 100644
--- a/op_analysis/gen_data.py
+++ b/op_analysis/gen_data.py
@@ -1,11 +1,10 @@
 import yaml
 import csv
 import torch
-import functorch
-import re
 import sys
 import os
 
+
 class CapturedOutput(object):
     """
     Class used to grab standard output.
@@ -60,6 +59,7 @@ def readOutput(self):
                 break
             self.capturedtext += char
 
+
 def get_ops_for_key(key):
     all_out = CapturedOutput()
     with all_out:
@@ -76,6 +76,7 @@ def get_ops_for_key(key):
         cleaned_ops.append(i[6:].strip())
     return set(cleaned_ops)
 
+
 def gen_data(special_op_lists, analysis_name):
     all_ops = get_ops_for_key(None)
     composite_ops = get_ops_for_key('CompositeImplicitAutograd')
@@ -83,7 +84,7 @@ def gen_data(special_op_lists, analysis_name):
 
     ops = yaml.load(open('../../pytorch/aten/src/ATen/native/native_functions.yaml', 'r').read(), Loader=yaml.CLoader)
 
-    annotated_ops = {a.strip(): b.strip() for a,b in list(csv.reader(open('annotated_ops.txt')))}
+    annotated_ops = {a.strip(): b.strip() for a, b in list(csv.reader(open('annotated_ops.txt')))}
     from collections import defaultdict
 
     uniq_ops = []
@@ -112,7 +113,6 @@ def gen_data(special_op_lists, analysis_name):
     def annotate_ops(ops, is_unique):
         categorization = defaultdict(int)
         for op in ops:
-            old_tcnt = sum(categorization.values())
             if op['name'][-1] == '_':
                 categorization['inplace'] += 1
                 op['meta'] = 'inplace'
@@ -141,11 +141,12 @@ def annotate_ops(ops, is_unique):
                 categorization['batch_norm'] += 1
                 op['meta'] = 'batch_norm'
                 continue
-            if 'Tensor' not in op['func'] or'Tensor' not in op['ret_type']:
+            if 'Tensor' not in op['func'] or 'Tensor' not in op['ret_type']:
                 categorization['non_tensor'] += 1
                 op['meta'] = 'non_tensor'
                 continue
-            if 'cudnn' in op['name'] or 'mkldnn' in op['name'] or 'miopen' in op['name'] or 'native' in op['name'] or 'thnn' in op['name'] or 'slow' in op['name']:
+            if 'cudnn' in op['name'] or 'mkldnn' in op['name'] or 'miopen' in op['name'] or \
+                    'native' in op['name'] or 'thnn' in op['name'] or 'slow' in op['name']:
                 categorization['backend'] += 1
                 op['meta'] = 'backend'
                 continue
@@ -157,19 +158,22 @@ def annotate_ops(ops, is_unique):
                 op['meta'] = 'core unknown'
         return categorization
 
-    # categorization = annotate_ops(uniq_ops, True)
-    categorization = annotate_ops(ops, False)
-
     with open(f"{analysis_name}", 'w') as f:
         for op in ops:
-            info = [op['full_name'], op['meta'], not (op['full_name'] in noncomposite_ops)] + [check(op) for check in special_op_lists]
+            info = [
+                op['full_name'], op['meta'], not (op['full_name'] in noncomposite_ops)
+            ] + [check(op) for check in special_op_lists]
             f.write(','.join([str(i) for i in info]) + '\n')
 
+
 def name_check(lst):
     return lambda x: x['name'] in lst
+
+
 def full_name_check(lst):
     return lambda x: x['full_name'] in lst
 
+
 # Generates batching rule data
 gen_data([full_name_check(get_ops_for_key('FuncTorchBatched'))], 'vmap')
 
@@ -179,4 +183,4 @@ def full_name_check(lst):
         opinfo_ops = [i.strip() for i in f.readlines()]
     with open('run_decompositions.txt', 'r') as f:
         decomposed_ops = [i.strip() for i in f.readlines()]
-    gen_data([name_check(opinfo_ops), name_check(decomposed_ops)], 'decompositions')
\ No newline at end of file
+    gen_data([name_check(opinfo_ops), name_check(decomposed_ops)], 'decompositions')
diff --git a/setup.cfg b/setup.cfg
new file mode 100644
index 000000000..5ff814b31
--- /dev/null
+++ b/setup.cfg
@@ -0,0 +1,17 @@
+[bdist_wheel]
+universal=1
+
+[metadata]
+license_file = LICENSE
+
+[pep8]
+max-line-length = 120
+
+[flake8]
+max-line-length = 120
+exclude = docs, benchmarks
+per-file-ignores =
+    __init__.py: F401
+
+[pydocstyle]
+select = D417 # Missing argument descriptions in the docstring


Test patch:
diff --git a/test/common_utils.py b/test/common_utils.py
index 776d1b515..a572ffcc8 100644
--- a/test/common_utils.py
+++ b/test/common_utils.py
@@ -17,11 +17,10 @@
 
 IS_FBCODE = os.getenv('FUNCTORCH_TEST_FBCODE') == '1'
 
+
 def loop(op, in_dims, out_dim, batch_size, *batched_args, **kwarg_values):
     outs = []
     for idx in range(batch_size):
-        idx_args = []
-        idx_kwargs = {}
         flat_args, args_spec = pytree.tree_flatten(batched_args)
         flat_dims, dims_spec = pytree.tree_flatten(in_dims)
         assert(args_spec == dims_spec)
@@ -56,6 +55,7 @@ def add_batch_dim(arg, bdim, batch_size=3):
             return (arg, None)
     for bdim in bdims:
         batch_choices = []
+
         def add_batch_choices(a):
             if isinstance(a, torch.Tensor):
                 batched_val = add_batch_dim(a, bdim, batch_size)
@@ -111,19 +111,25 @@ def f(x, *args, **kwargs):
             loop_out = pytree.tree_map(lambda v: torch.ones(2, *v.shape, dtype=v.dtype, device=v.device) + v, loop_out)
         else:
             loop_out = None
-        batched_out = vmap(vmap(f, in_dims=vmap1_dims), in_dims=vmap2_dims)(torch.ones(2), *batched_args, **kwarg_values)
+        batched_out = vmap(vmap(f, in_dims=vmap1_dims), in_dims=vmap2_dims)(
+            torch.ones(2), *batched_args, **kwarg_values)
         yield (loop_out, batched_out)
 
+
 def opinfo_in_dict(opinfo, d):
     return (opinfo.name in d) or (f'{opinfo.name}.{opinfo.variant_test_name}' in d)
 
+
 def xfail(op_name, variant_name=None, *, device_type=None, dtypes=None):
     return (op_name, variant_name, device_type, dtypes, True)
 
 # TODO: this doesn't work in python < 3.8
+
+
 def skip(op_name, variant_name=None, *, device_type=None, dtypes=None):
     return (op_name, variant_name, device_type, dtypes, False)
 
+
 def skipOps(test_case_name, base_test_name, to_skip):
     all_opinfos = functorch_lagging_op_db + additional_op_db
     for xfail in to_skip:
@@ -153,6 +159,7 @@ def wrapped(fn):
         return fn
     return wrapped
 
+
 class DisableVmapFallback:
     def __enter__(self):
         self.prev_state = functorch._C._is_vmap_fallback_enabled()
@@ -161,11 +168,12 @@ def __enter__(self):
     def __exit__(self, *ignored):
         functorch._C._set_vmap_fallback_enabled(self.prev_state)
 
+
 def check_vmap_fallback(test_case, thunk, opinfo, dry_run=False):
     try:
         with DisableVmapFallback():
             thunk()
-    except:
+    except Exception:
         if not dry_run:
             raise
         if opinfo.variant_test_name:
diff --git a/test/discover_coverage.py b/test/discover_coverage.py
index ccccdfb14..f02b17ba3 100644
--- a/test/discover_coverage.py
+++ b/test/discover_coverage.py
@@ -2,13 +2,12 @@
 import copy
 from torch.testing._internal.common_methods_invocations import op_db
 from enum import Enum
-from functorch_lagging_op_db import functorch_lagging_op_db
 import functorch._src.top_operators_github_usage as top_ops
 import pprint
 
 # Importing these files make modifications to the op_db that we need
-import test_ops
-import test_vmap
+import test_ops  # noqa: F401
+import test_vmap  # noqa: F401
 
 all_overridable = list(torch.overrides.get_testing_overrides().keys())
 
@@ -22,6 +21,8 @@
 ]
 
 # torch.abs, Tensor.abs, Tensor.abs_ are all considered to be different
+
+
 def get_public_overridable_apis(pytorch_root='/raid/rzou/pt/quick'):
     results = {}
     all_overridable_apis = set(torch.overrides.get_testing_overrides().keys())
@@ -41,6 +42,7 @@ def get_public_overridable_apis(pytorch_root='/raid/rzou/pt/quick'):
                 results[f'{module_name}.{line}'] = api
     return results
 
+
 denylist = {
     'torch.Tensor.data_ptr',
     'torch.Tensor.dim',
@@ -74,6 +76,7 @@ def get_public_overridable_apis(pytorch_root='/raid/rzou/pt/quick'):
     'torch.Tensor.storage_type',
 }
 
+
 def get_method_only_ops_we_care_about():
     apis = get_public_overridable_apis()
     result = []
@@ -91,6 +94,8 @@ def get_method_only_ops_we_care_about():
     return result
 
 # Deduplicates torch.abs and Tensor.abs
+
+
 def get_public_overridable_ops():
     results = get_public_overridable_apis()
     cpy = copy.deepcopy(results)
@@ -102,6 +107,7 @@ def get_public_overridable_ops():
             del results[key]
     return results
 
+
 def get_public_overridable_outplace_ops():
     results = get_public_overridable_ops()
     cpy = copy.deepcopy(results)
@@ -111,6 +117,7 @@ def get_public_overridable_outplace_ops():
             del results[key]
     return results
 
+
 def get_public_overridable_outplace_we_care_about():
     results = get_public_overridable_outplace_ops()
     cpy = copy.deepcopy(results)
@@ -128,6 +135,8 @@ def get_public_overridable_outplace_we_care_about():
     return results
 
 # e.g. nn.functional.softmax
+
+
 def get_op(dotted_name):
     names = dotted_name.split('.')
     mod = torch
@@ -138,6 +147,8 @@ def get_op(dotted_name):
     return mod
 
 # Maps function -> [OpInfo]
+
+
 def get_ops_covered_by_opinfos():
     ops = {}
 
@@ -159,12 +170,14 @@ def safe_append(dct, key, val):
             safe_append(ops, alias.op, opinfo)
     return ops
 
+
 factory_fns = {
     'tensor', 'zeros', 'ones', 'randn', 'arange', 'rand', 'empty', 'randperm',
     'linspace', 'logspace', 'hann_window', 'full', 'eye', 'blackman_window',
     'barlett_window', 'randint', 'range', 'arange',
 }
 
+
 def get_top_ops(torch_threshold, nn_fn_threshold):
     denylist = set({
         # These are either not real "operators", factory functions
@@ -207,7 +220,7 @@ def get_top_ops(torch_threshold, nn_fn_threshold):
         'nn.functional.relu_',
         'nn.functional.boolean_dispatch',
         'nn.functional.assert_int_or_pair',
-        'fft', # is namespace
+        'fft',  # is namespace
     })
 
     torch_ops = [op[0] for op in top_ops.top_torch]
@@ -218,6 +231,7 @@ def get_top_ops(torch_threshold, nn_fn_threshold):
     ops = torch_ops[:torch_threshold] + nn_fn_ops[:nn_fn_threshold]
     return ops
 
+
 def get_top_ops_not_covered_by_opinfo(torch_threshold=0, nn_fn_threshold=0):
     ops = get_top_ops(torch_threshold, nn_fn_threshold)
 
@@ -232,6 +246,7 @@ def get_top_ops_not_covered_by_opinfo(torch_threshold=0, nn_fn_threshold=0):
     result = [op for op in result if op not in factory_fns]
     return result
 
+
 def get_covered_ops(ops_list, invert=False):
     ops_covered_by_opinfo = get_ops_covered_by_opinfos()
     overridable_outplace_ops = ops_list
@@ -244,10 +259,12 @@ def get_covered_ops(ops_list, invert=False):
             results[key] = op
     return results
 
+
 class Status(Enum):
     Correct = 0
     Fast = 1
 
+
 tests = {
     'test_vmap_exhaustive',
     'test_op_has_batch_rule',
@@ -258,6 +275,7 @@ class Status(Enum):
     'test_vmapjvp',
 }
 
+
 def get_statuses(for_subset=None, invert=False):
     overridable_outplace_we_care_about = get_public_overridable_outplace_we_care_about()
     if for_subset is not None:
@@ -269,7 +287,7 @@ def get_statuses(for_subset=None, invert=False):
         }
     op_to_opinfo = get_ops_covered_by_opinfos()
     result = {}
-    x = get_covered_ops(overridable_outplace_we_care_about)
+    _ = get_covered_ops(overridable_outplace_we_care_about)
 
     def get_covered_tests(op):
         opinfos = op_to_opinfo[op]
@@ -296,6 +314,7 @@ def get_all_aliases(op):
         result[name] = failed_tests if invert else successful_tests
     return result
 
+
 def transpose_statuses(for_subset=None, invert=False):
     statuses = get_statuses(for_subset, invert=invert)
     result = {}
@@ -306,6 +325,7 @@ def transpose_statuses(for_subset=None, invert=False):
             result[test].add(op)
     return result
 
+
 overridable_apis = get_public_overridable_apis()
 
 overridable_ops = get_public_overridable_ops()
@@ -348,21 +368,23 @@ def transpose_statuses(for_subset=None, invert=False):
 # for op in top_ops_not_covered_by_opinfo:
 #     print(f'{op}, {top_ops.usage_count[op]}')
 
-#print("top ops not covered by opinfo: ")
-#top_ops_not_covered_by_opinfo = get_top_ops_not_covered_by_opinfo(220, 92)
-#for op in top_ops_not_covered_by_opinfo:
+# print("top ops not covered by opinfo: ")
+# top_ops_not_covered_by_opinfo = get_top_ops_not_covered_by_opinfo(220, 92)
+# for op in top_ops_not_covered_by_opinfo:
 #    print(f'{op}, {top_ops.usage_count[op]}')
 
-#print("top ops not covered by opinfo: ")
+# print("top ops not covered by opinfo: ")
 # top_ops_not_covered_by_opinfo = get_top_ops_not_covered_by_opinfo(999, 999)
 # for op in top_ops_not_covered_by_opinfo:
 #     print(f'{op}, {top_ops.usage_count[op]}')
 
+
 def remove_from_set(parent, to_remove):
     for to_remove_elt in to_remove:
         if to_remove_elt in parent:
             parent.remove(to_remove_elt)
 
+
 def print_coverage_info(th=100, nn=25):
     print('=' * 80)
     print(f"top {th}, {nn} coverage")
@@ -371,20 +393,20 @@ def print_coverage_info(th=100, nn=25):
 
     # testing problems
     exemptions = {
-        'torch.nn.functional.dropout', # randomness
+        'torch.nn.functional.dropout',  # randomness
     }
 
     # Allowed exemptions
     vmap_exemptions = {
-        'torch.randn_like', # randomness
-        'torch.rand_like', # randomness
-        'torch.allclose', # number output
-        'torch.unique', # dynamic
-        'torch.nonzero', # dynamic
-        'torch.masked_select', # dynamic
-        'torch.prod', # dynamic (backward)
-        'torch.norm', # norm with nuc is not commonly used; we support the other cases.
-        'torch.svd', # There isn't a bug, it is just nondeterministic so we can't test it.
+        'torch.randn_like',  # randomness
+        'torch.rand_like',  # randomness
+        'torch.allclose',  # number output
+        'torch.unique',  # dynamic
+        'torch.nonzero',  # dynamic
+        'torch.masked_select',  # dynamic
+        'torch.prod',  # dynamic (backward)
+        'torch.norm',  # norm with nuc is not commonly used; we support the other cases.
+        'torch.svd',  # There isn't a bug, it is just nondeterministic so we can't test it.
     }
     remove_from_set(statuses['test_vmap_exhaustive'], vmap_exemptions)
     remove_from_set(statuses['test_vmapvjp'], vmap_exemptions)
@@ -407,6 +429,7 @@ def print_coverage_info(th=100, nn=25):
 
     pprint.pprint(statuses)
 
+
 print_coverage_info(100, 25)
 # print_coverage_info(200, 50)
 
diff --git a/test/functorch_additional_op_db.py b/test/functorch_additional_op_db.py
index 625b626e7..17f44c730 100644
--- a/test/functorch_additional_op_db.py
+++ b/test/functorch_additional_op_db.py
@@ -1,45 +1,13 @@
-from functools import wraps, partial
-from itertools import product, chain
+from functools import partial
 import itertools
-import collections
-import copy
-import operator
-import random
-import numbers
-import unittest
 
 import torch
-import numpy as np
-from torch._six import inf
-import collections.abc
-
-from typing import Any, List, Sequence, Tuple, Union
 
 from torch.testing import \
-    (make_non_contiguous, floating_types, floating_types_and, complex_types,
-     floating_and_complex_types, floating_and_complex_types_and,
-     all_types_and_complex_and, all_types_and, all_types_and_complex,
-     integral_types_and, all_types)
-# from .._core import _dispatch_dtypes
-from torch.testing._internal.common_device_type import \
-    (skipIf, skipCUDAIfNoMagma, skipCUDAIfNoMagmaAndNoCusolver, skipCUDAIfNoCusolver,
-     skipCPUIfNoLapack, skipCPUIfNoFFT, skipCUDAIfRocm, precisionOverride, toleranceOverride, tol)
-from torch.testing._internal.common_cuda import CUDA11OrLater, SM53OrLater, SM60OrLater
-from torch.testing._internal.common_utils import \
-    (is_iterable_of_tensors,
-     random_symmetric_matrix, random_symmetric_psd_matrix,
-     make_fullrank_matrices_with_distinct_singular_values,
-     random_symmetric_pd_matrix, make_symmetric_matrices,
-     make_symmetric_pd_matrices,
-     random_fullrank_matrix_distinct_singular_value,
-     TEST_WITH_ROCM, IS_WINDOWS, IS_MACOS, make_tensor, TEST_SCIPY,
-     torch_to_numpy_dtype_dict, TEST_WITH_ASAN,
-     GRADCHECK_NONDET_TOL,)
-import torch.testing._internal.opinfo_helper as opinfo_helper
-from torch.testing._internal.common_methods_invocations import (
-    OpInfo, DecorateInfo, SampleInput, sample_inputs_softshrink_hardshrink_hardtanh,
-    sample_inputs_softmax_variant, S
-)
+    (floating_types, floating_types_and, floating_and_complex_types_and,
+     all_types_and_complex_and)
+from torch.testing._internal.common_utils import make_tensor
+from torch.testing._internal.common_methods_invocations import OpInfo, SampleInput
 
 # List of OpInfos that aren't in PyTorch Core yet.
 # They are here because we wanted a fast way of writing OpInfos and may not be
@@ -49,10 +17,12 @@
 additional_op_db = []
 
 # https://github.com/pytorch/pytorch/pull/61971
+
+
 def sample_inputs_linear(has_bias, self, device, dtype, requires_grad):
     features_options = [[3, 4], [128, 128]]
     batch_options = [
-        [], # no batch
+        [],  # no batch
         [64],
         [5, 7],
     ]
@@ -75,6 +45,7 @@ def sample_inputs_linear(has_bias, self, device, dtype, requires_grad):
         sample_inputs.append(SampleInput(input_tensor, args=(weight, bias)))
     return sample_inputs
 
+
 additional_op_db.extend([
     OpInfo('nn.functional.linear',
            aten_name='linear',
@@ -96,6 +67,7 @@ def sample_inputs_linear(has_bias, self, device, dtype, requires_grad):
 
 # https://github.com/pytorch/pytorch/pull/61068
 
+
 def sample_inputs_conv2d(has_bias, self, device, dtype, requires_grad, extra_args=(), groups=1):
     in_ch, out_ch = 6, 4
     inp = make_tensor((2, in_ch * groups, 7, 5), device=device, dtype=dtype,
@@ -108,6 +80,7 @@ def sample_inputs_conv2d(has_bias, self, device, dtype, requires_grad, extra_arg
                            requires_grad=requires_grad, low=-1, high=1)
     return [SampleInput(inp, args=((weight, bias) + extra_args))]
 
+
 additional_op_db.extend([
     OpInfo('nn.functional.conv2d',
            aten_name="conv2d",
@@ -191,6 +164,7 @@ def sample_inputs_conv2d(has_bias, self, device, dtype, requires_grad, extra_arg
            supports_out=False),
 ])
 
+
 def sample_inputs_cross_entropy(self, device, dtype, requires_grad, reduction):
     N = 2
     C = 10
@@ -214,6 +188,7 @@ def sample_inputs_cross_entropy(self, device, dtype, requires_grad, reduction):
     ])
     return sample_inputs
 
+
 for reduction in ['mean', 'sum', 'none']:
     additional_op_db.append(
         OpInfo('nn.functional.cross_entropy',
@@ -225,6 +200,7 @@ def sample_inputs_cross_entropy(self, device, dtype, requires_grad, reduction):
                dtypesIfCUDA=floating_types_and(torch.half, torch.bfloat16),
                supports_out=True))
 
+
 # TODO: split embedding in pytorch core
 def sample_inputs_embedding(op_info, device, dtype, requires_grad, **kwargs):
     def make_input(shape):
@@ -273,6 +249,7 @@ def generator():
 
     return list(generator())
 
+
 additional_op_db.append(
     OpInfo(
         "nn.functional.embedding",
diff --git a/test/test_compile_cache.py b/test/test_compile_cache.py
index 0b1c1e303..7abf63d24 100644
--- a/test/test_compile_cache.py
+++ b/test/test_compile_cache.py
@@ -1,20 +1,11 @@
 import torch
-import unittest
 
-import numpy as np
-
-import torch
-from torch import fx
 import functorch
-from torch.testing._internal.common_utils import run_tests
-from torch.testing._internal.jit_utils import JitTestCase
-import unittest
+from torch.testing._internal.common_utils import run_tests, TestCase
 
 from functorch.compile import memory_efficient_pointwise_fusion
 from functorch.compile import compiled_function
 
-from torch.testing._internal.common_utils import TestCase, run_tests
-
 
 class TestCompileCache(TestCase):
     def test_recompilation_on_broadcast(self):
diff --git a/test/test_eager_transforms.py b/test/test_eager_transforms.py
index 8d7e929d3..6443344d3 100644
--- a/test/test_eager_transforms.py
+++ b/test/test_eager_transforms.py
@@ -15,8 +15,7 @@
 import warnings
 import math
 from typing import Callable, Type
-from torch.testing._internal.common_device_type import instantiate_device_type_tests, \
-    skipCUDAIfNoMagma, onlyCPU
+from torch.testing._internal.common_device_type import instantiate_device_type_tests, onlyCPU
 from functools import partial
 
 import functorch
@@ -38,7 +37,7 @@
 
 USE_TORCHVISION = False
 try:
-    import torchvision
+    import torchvision  # noqa: F401
     USE_TORCHVISION = True
 except ImportError:
     warnings.warn("Couldn't import torchvision. Some of our tests use it, try "
@@ -47,6 +46,8 @@
                   UserWarning)
 
 # TestCase for _argnums_partial, an important helper funciton
+
+
 class TestArgnumsPartial(TestCase):
     def test_invalid_argnum_type(self):
         x = torch.randn(3)
@@ -59,6 +60,7 @@ def test_invalid_argnum_type(self):
             _argnums_partial(torch.sin, args, (0.0,))
 
         args = (0.1, 1.1, 2.1, 3.1, 4.1)
+
         def f(a, b, c, d, e):
             return a
         with self.assertRaisesRegex(RuntimeError, "must be int"):
@@ -90,6 +92,7 @@ def test_duplicate_argnums(self):
 
     def test_flat_args_with_positive_int_argnum(self):
         args = (0.1, 1.1, 2.1, 3.1, 4.1)
+
         def f(a, b, c, d, e):
             return a
 
@@ -103,6 +106,7 @@ def f(a, b, c, d, e):
 
     def test_flat_args_with_negative_int_argnum(self):
         args = (0.1, 1.1, 2.1, 3.1, 4.1)
+
         def f(a, b, c, d, e):
             return a
 
@@ -117,6 +121,7 @@ def f(a, b, c, d, e):
 
     def test_flat_args_with_tuple_argnum(self):
         args = (0.1, 1.1, 2.1, 3.1, 4.1)
+
         def f(a, b, c, d, e):
             return a
 
@@ -130,6 +135,7 @@ def f(a, b, c, d, e):
 
     def test_pytree_args(self):
         args = ((0.1, 1.1), 2.0, [3.1])
+
         def f(a, b, c):
             return a[0] + a[1] + b + c[0]
 
@@ -153,6 +159,7 @@ def f(a, b, c):
 
     def test_argnums_reorders(self):
         args = ((0.1, 1.1, 2.1), 3.1, 4.1)
+
         def f(a, b, c):
             return a[0] + a[1] + a[2] + b + c
 
@@ -163,6 +170,7 @@ def f(a, b, c):
 
     def test_function_with_default_args(self):
         args = ((0.1, 1.1, 2.1), 3.1)
+
         def f(a, b, c=4.1):
             return a[0] + a[1] + a[2] + b + c
 
@@ -177,6 +185,7 @@ def f(a, b, c=4.1):
         self.assertEqual(res, args[-1:])
         self.assertEqual(f_new(*res), expected)
 
+
 class TestGradTransform(TestCase):
     def test_primitive(self, device):
         x = torch.randn([], device=device)
@@ -313,23 +322,25 @@ def test_nesting_simple(self, device):
     def test_escaped_wrappers_are_marked_as_dead(self, device):
         x = torch.randn([], device=device)
         escaped = []
+
         def foo(x):
             y = x.sin()
             escaped.append(y)
             return y
 
-        result = grad(foo)(x)
+        grad(foo)(x)
         self.assertEqual(functorch._C.dlevel(escaped[0]), -1)
 
     def test_escaped_wrappers_are_ignored(self, device):
         x = torch.randn([], device=device)
         escaped = []
+
         def foo(x):
             y = x.sin()
             escaped.append(y)
             return y
 
-        result = grad(foo)(x)
+        grad(foo)(x)
 
         something = escaped[0].sum()
         self.assertEqual(functorch._C.dlevel(something), 0)
@@ -352,6 +363,7 @@ def f(x):
 
     def test_conj_bit(self):
         x = torch.tensor(1+1j)
+
         def foo(x):
             assert not x.is_conj()
             y = x.conj()
@@ -489,6 +501,7 @@ def test_negative_argnums(self, device):
 
     def test_grad_pytree_inputs(self, device):
         x = torch.randn([], device=device)
+
         def f(a, b):
             x, y = a
             return 1 * x + 2 * y + 3 * b['foo']
@@ -511,7 +524,7 @@ def f(a, b):
     def test_zero_grad(self, device):
         def f(x):
             return (x['a']**2.0).sum()
-        inps = ({'a':torch.randn(10, device=device) + 3, 'b':torch.randn(10, device=device)})
+        inps = ({'a': torch.randn(10, device=device) + 3, 'b': torch.randn(10, device=device)})
         grads = grad(f)(inps)
         self.assertNotEqual(grads['a'].sum(), 0.0)
         self.assertEqual(grads['b'].sum(), 0.0)
@@ -856,6 +869,7 @@ def h(x):
         z, = torch.autograd.grad(y, x)
         self.assertEqual(z, 2)
 
+
 class TestVmapOfGrad(TestCase):
     def test_per_sample_grads_inplace_view(self, device):
         def compute_loss(weight, x, t):
@@ -977,10 +991,12 @@ def foo(x, v):
         output.backward(v)
         self.assertEqual(result, x.grad)
 
+
 jacrev_and_jacfwd = parametrize("jacapi", [subtest(jacrev, name='jacrev'), subtest(jacfwd, name='jacfwd')])
 
 FIXME_jacrev_only = parametrize("jacapi", [subtest(jacrev, name='jacrev')])
 
+
 class TestJac(TestCase):
     @jacrev_and_jacfwd
     def test_simple(self, device, jacapi):
@@ -1124,7 +1140,7 @@ def f(a, b, c):
         expected = ((torch.tensor(1., device=device), torch.tensor(2., device=device)),)
         self.assertEqual(result, expected)
 
-        result = jacapi(f )(*args)
+        result = jacapi(f)(*args)
         expected = (torch.tensor(1., device=device), torch.tensor(2., device=device))
         self.assertEqual(result, expected)
 
@@ -1248,6 +1264,7 @@ def f(x):
     def test_empty_output(self, device, jacapi):
         x = torch.randn(3, device=device)
         y = torch.randn(3, device=device)
+
         def f(x, y):
             return ()
 
@@ -1297,40 +1314,40 @@ def f(x, y):
     def test_empty_argnums(self, device, jacapi):
         x = torch.randn(3, device=device)
         with self.assertRaisesRegex(RuntimeError, "must be non-empty"):
-            z = jacapi(torch.sin, argnums=())(x)
+            jacapi(torch.sin, argnums=())(x)
 
     @jacrev_and_jacfwd
     def test_out_of_bounds_argnums(self, device, jacapi):
         x = torch.randn(3, device=device)
         with self.assertRaisesRegex(RuntimeError, "only 1 positional inputs"):
-            z = jacapi(torch.sin, argnums=2)(x)
+            jacapi(torch.sin, argnums=2)(x)
 
     @jacrev_and_jacfwd
     def test_negative_argnums(self, device, jacapi):
         x = torch.randn(3, device=device)
         with self.assertRaisesRegex(RuntimeError, "only 1 positional inputs"):
-            z = jacapi(torch.sin, argnums=-2)(x)
+            jacapi(torch.sin, argnums=-2)(x)
 
     @jacrev_and_jacfwd
     def test_repeated_argnums(self, device, jacapi):
         x = torch.randn(3, device=device)
         with self.assertRaisesRegex(RuntimeError, "must be unique"):
-            z = jacapi(torch.sin, argnums=(0, 0))(x)
+            jacapi(torch.sin, argnums=(0, 0))(x)
 
     @jacrev_and_jacfwd
     def test_float_argnums(self, device, jacapi):
         x = torch.randn(3, device=device)
         with self.assertRaisesRegex(RuntimeError, "must be int or Tuple"):
-            z = jacapi(torch.sin, argnums=0.0)(x)
+            jacapi(torch.sin, argnums=0.0)(x)
         with self.assertRaisesRegex(RuntimeError, "must be int"):
-            z = jacapi(torch.multiply, argnums=(1, 0.0))(x, x)
+            jacapi(torch.multiply, argnums=(1, 0.0))(x, x)
 
     def test_hessian_simple(self, device):
         def f(x):
             return x.sin()
 
         x = torch.randn(3, device=device)
-        result = hessian(f)(x)
+        hessian(f)(x)
 
     def _test_against_reference(self, f, inputs, jacapi):
         def foo(inputs):
@@ -1409,6 +1426,7 @@ def f(x, y):
         y = torch.randn(3)
         self._test_against_reference(f, (x, y), jacapi)
 
+
 class TestHessian(TestCase):
     def _test_against_reference(self, f, inputs):
         def foo(inputs):
@@ -1451,6 +1469,7 @@ def f(x, y):
         y = torch.randn(3, device=device)
         self._test_against_reference(f, (x, y))
 
+
 class TestJvp(TestCase):
     def test_inplace_on_captures(self, device):
         x = torch.tensor([1., 2., 3.], device=device)
@@ -1627,6 +1646,7 @@ def h(x):
         with self.assertRaisesRegex(RuntimeError, "a Tensor or Tensors"):
             jvp(h, (x,), (t,))
 
+
 class TestCustomFunction(TestCase):
     @onlyCPU
     def test_basic(self, device):
@@ -1660,6 +1680,7 @@ def filter_fn(args):
 
         assert torch.allclose(x.grad, 3 * x.cos())
 
+
 class TestComposability(TestCase):
     def test_grad_grad(self, device):
         x = torch.randn([], device=device)
@@ -1767,7 +1788,6 @@ def mse_loss(x, y):
 
         net, params = make_functional(ThreeLayerNet().to(device))
         K = 20
-        losses = []
         num_tasks = 4
         alpha = 0.1
 
@@ -1778,6 +1798,7 @@ def sample_tasks(outer_batch_size, inner_batch_size):
             for _ in range(outer_batch_size):
                 As.append(np.random.uniform(low=0.1, high=.5))
                 phases.append(np.random.uniform(low=0., high=np.pi))
+
             def get_batch():
                 xs, ys = [], []
                 for A, phase in zip(As, phases):
@@ -1810,8 +1831,7 @@ def inner_loss(params, x1, y1):
         task = sample_tasks(num_tasks, K)
 
         # Compute with vmap+grad
-        inner_losses = vmap(partial(get_loss_for_task, True))\
-                            (task[0], task[1], task[2], task[3])
+        inner_losses = vmap(partial(get_loss_for_task, True))(task[0], task[1], task[2], task[3])
         loss2 = sum(inner_losses)/len(inner_losses)
         result_grads = torch.autograd.grad(loss2, params)
 
@@ -1835,6 +1855,7 @@ def test_maml_omniglot(self, device):
         n_way = 5
         n_inner_iter = 2
         num_tasks = 2
+
         class Flatten(nn.Module):
             def forward(self, input):
                 return input.view(input.size(0), -1)
@@ -1886,7 +1907,7 @@ def compute_loss(new_params, buffers, x, y):
         # Get some sample inputs...
         x_spt = torch.randn(num_tasks, 25, 1, 28, 28, dtype=dtype, device=device)
         y_spt = torch.randint(0, 5, (num_tasks, 25), device=device)
-        x_qry = torch.randn(num_tasks, 75, 1, 28, 28, dtype=dtype,device=device)
+        x_qry = torch.randn(num_tasks, 75, 1, 28, 28, dtype=dtype, device=device)
         y_qry = torch.randint(0, 5, (num_tasks, 75), device=device)
 
         # compute with vmap + grad
@@ -2112,6 +2133,7 @@ def compute_loss(weights, image, target):
 
         self.assertEqual(result_grads, expected_grads)
 
+
 only_for = ("cpu", "cuda")
 instantiate_device_type_tests(
     TestGradTransform,
diff --git a/test/test_minifier.py b/test/test_minifier.py
index bb0654f2a..09454c2de 100644
--- a/test/test_minifier.py
+++ b/test/test_minifier.py
@@ -1,9 +1,7 @@
 import torch
-from torch.nn import functional as F
 from functorch.compile import minimizer
 from functorch import make_fx
 from torch.testing._internal.common_utils import TestCase, run_tests
-from typing import Callable
 
 
 class TestMinifier(TestCase):
@@ -21,7 +19,7 @@ def pass_checker(fx_g, inps):
 
         min_f, inps = minimizer(failing_f, inps, pass_checker)
         assert len(min_f.graph.nodes) == 4
-        assert len(inps) == 2 
+        assert len(inps) == 2
 
     def test_has_add_mul(self):
         def failing_f(x):
@@ -45,8 +43,8 @@ def pass_checker(fx_g, inps):
 
         min_f, inps = minimizer(failing_f, inps, pass_checker)
         assert len(min_f.graph.nodes) == 3
-        assert len(inps) == 1 
+        assert len(inps) == 1
 
 
 if __name__ == "__main__":
-    run_tests()
\ No newline at end of file
+    run_tests()
diff --git a/test/test_operator_authoring.py b/test/test_operator_authoring.py
index cc1bf0507..f34ac9025 100644
--- a/test/test_operator_authoring.py
+++ b/test/test_operator_authoring.py
@@ -1,20 +1,16 @@
 import torch
 import unittest
 
-import numpy as np
-
-import torch
 from torch import fx
 from functorch.compile import pointwise_operator
 from torch.testing._internal.common_utils import run_tests
 from torch.testing._internal.jit_utils import JitTestCase
-import unittest
 
 LLVM_ENABLED = torch._C._llvm_enabled()
 HAS_CUDA = torch.cuda.is_available()
 HAS_SYMPY = False
 try:
-    import sympy
+    import sympy  # noqa: F401
 
     HAS_SYMPY = True
 except ImportError:
@@ -186,7 +182,7 @@ def test_unary_ops(self):
         ]
 
         for unary_op in unary_operators:
-            fn = lambda x: unary_op(x)
+            def fn(x): return unary_op(x)
             pointwise_fn = pointwise_operator(fn)
             a = torch.rand(2, 3)
             ref = fn(a)
@@ -218,7 +214,7 @@ def test_binary_ops(self):
             torch.ops.aten.atan2,
         ]
         for binary_op in binary_operators:
-            fn = lambda x, y: binary_op(x, y)
+            def fn(x, y): return binary_op(x, y)
             pointwise_fn = pointwise_operator(fn)
             a = torch.rand(2, 3)
             b = torch.rand(2, 3)
diff --git a/test/test_ops.py b/test/test_ops.py
index 25317eaed..927ed62a4 100644
--- a/test/test_ops.py
+++ b/test/test_ops.py
@@ -6,30 +6,23 @@
 
 from torch.testing._internal.common_utils import TestCase, run_tests, is_iterable_of_tensors
 import torch
-import torch.nn.functional as F
 from torch import Tensor
 import functools
-import itertools
-import copy
-import warnings
 import unittest
-from torch.testing._internal.common_device_type import instantiate_device_type_tests, \
-    skipCUDAIfNoMagma
-from torch.testing._internal.common_device_type import ops, onlyCPU
-from torch.testing._internal.common_dtype import floating_types_and, integral_types
+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_dtype import integral_types
 from functorch_lagging_op_db import functorch_lagging_op_db
 from functorch_additional_op_db import additional_op_db
 from common_utils import (
     get_fallback_and_vmap_exhaustive,
     get_exhaustive_batched_inputs,
-    opinfo_in_dict,
     xfail,
     skip,
     skipOps,
     check_vmap_fallback,
     IS_FBCODE,
 )
-import types
 from torch.utils._pytree import tree_flatten, tree_unflatten, tree_map
 from functorch import grad, vjp, vmap
 import torch.autograd.forward_ad as fwAD
@@ -38,6 +31,8 @@
 aten = torch.ops.aten
 
 # Version of autograd.grad that handles outputs that don't depend on inputs
+
+
 def _autograd_grad(outputs, inputs, grad_outputs=None, retain_graph=False, create_graph=True):
     inputs, inputs_spec = tree_flatten(inputs)
     result = [torch.zeros_like(inp) for inp in inputs]
@@ -108,9 +103,13 @@ def wrapped(*primals):
         return result
     return wrapped, primals
 
+
 def normalize_op_input_output(f, sample, requires_grad=True):
     args = tuple([sample.input] + list(sample.args))
-    return normalize_op_input_output2(f, args, sample.kwargs, sample.output_process_fn_grad, requires_grad=requires_grad)
+    return normalize_op_input_output2(
+        f, args, sample.kwargs, sample.output_process_fn_grad, requires_grad=requires_grad
+    )
+
 
 def ref_vjp(f, *primals):
     result = f(*primals)
@@ -120,6 +119,7 @@ def wrapped(cotangents):
 
     return result, wrapped
 
+
 def ref_jvp(f, primals, tangents):
     with fwAD.dual_level():
         duals = tuple(fwAD.make_dual(p, t) for p, t in zip(primals, tangents))
@@ -130,6 +130,8 @@ def ref_jvp(f, primals, tangents):
 
 # Returns a new function g(*args, *cotangents) that computes vjps and
 # sample (*args, *cotangents)
+
+
 def get_vjpfull_variant(f, sample):
     fn, primals = normalize_op_input_output(f, sample)
     result = fn(*primals)
@@ -150,6 +152,7 @@ def wrapped(*args):
 
     return wrapped, args
 
+
 def get_jvp_variant(f, sample):
     # We want this higher-order variant of jvp, so that it can
     # be used to wrap vmap
@@ -171,6 +174,7 @@ def wrapped(*args):
 
     return wrapped, tangents
 
+
 def get_jvp_variant_primals_tangents(f, sample):
     # We want this higher-order variant of jvp, so that it can
     # be used to wrap vmap
@@ -193,6 +197,7 @@ def wrapped(*args):
 
     return wrapped, primals + tangents
 
+
 def is_inplace(op, variant):
     if hasattr(variant, "__wrapped__"):
         return variant.__wrapped__ is op.get_inplace()
@@ -212,6 +217,7 @@ def is_inplace(op, variant):
     xfail('nn.functional.fractional_max_pool2d'),
 }
 
+
 class TestOperators(TestCase):
     @ops(functorch_lagging_op_db + additional_op_db, allowed_dtypes=(torch.float,))
     @skipOps('TestOperators', 'test_grad', vjp_fail)
@@ -317,7 +323,7 @@ def test_jvp(self, device, dtype, op):
 
     @ops(functorch_lagging_op_db + additional_op_db, allowed_dtypes=(torch.float,))
     @skipOps('TestOperators', 'test_vjp', vjp_fail.union({
-	skip('nn.functional.conv_transpose3d', device_type='cuda'),  # numerical precision
+        skip('nn.functional.conv_transpose3d', device_type='cuda'),  # numerical precision
     }))
     def test_vjp(self, device, dtype, op):
         if not op.supports_autograd:
@@ -477,13 +483,14 @@ def vjp_of_vjp(*args_and_cotangents):
         xfail('fft.rfft2'),
         xfail('lu'),
         skip('qr'),  # Nondetermistic
-        xfail('_masked.prod'), # calls aten::item
+        xfail('_masked.prod'),  # calls aten::item
         xfail('stft'),
         xfail('nn.functional.glu'),
         xfail('nn.functional.fractional_max_pool3d'),
         xfail('as_strided'),
         xfail('nn.functional.fractional_max_pool2d'),
     })
+
     @ops(functorch_lagging_op_db + additional_op_db, allowed_dtypes=(torch.float,))
     @skipOps('TestOperators', 'test_vmapvjp', vmapvjp_fail)
     def test_vmapvjp(self, device, dtype, op):
@@ -541,7 +548,8 @@ def test_vmapvjp(self, device, dtype, op):
         xfail('max', 'binary'),
 
         # Apprently these support forward AD, but we get "Trying to use forward AD..."
-        # These are cases where OpInfo has supports_forward_ad=True, but disables the test
+        # These are cases where OpInfo has supports_forward_ad=True, but disables
+        # the test
         xfail('var_mean'),
         xfail('std_mean'),
         xfail('linalg.eigvalsh'),
@@ -555,7 +563,8 @@ def test_vmapvjp(self, device, dtype, op):
         # See https://github.com/pytorch/pytorch/issues/66357
         xfail('nn.functional.pad', 'circular'),
 
-        # RuntimeError: expand: the number of sizes provided (1) must be greater or equal to the number of dimensions in the tensor (2)
+        # RuntimeError: expand: the number of sizes provided (1) must be greater or
+        # equal to the number of dimensions in the tensor (2)
         xfail('nanquantile'),
         xfail('quantile'),
 
@@ -839,9 +848,9 @@ def test():
         xfail('double', 'channels_last'),
         xfail('masked_select'),
         xfail('nn.functional.fractional_max_pool3d'),
-	xfail('nn.functional.glu'),
-	xfail('as_strided'),
-	xfail('nn.functional.fractional_max_pool2d'),
+        xfail('nn.functional.glu'),
+        xfail('as_strided'),
+        xfail('nn.functional.fractional_max_pool2d'),
     }))
     def test_vjpvmap(self, device, dtype, op):
         # NB: there is no vjpvmap_has_batch_rule test because that is almost
@@ -870,7 +879,7 @@ def test_vjpvmap(self, device, dtype, op):
             for batched_args, in_dims, kwargs in get_exhaustive_batched_inputs(args, kwargs):
                 vmapped_op = vmap(op, in_dims)
                 fn, primals = normalize_op_input_output2(vmapped_op, batched_args, kwargs,
-                                                    sample.output_process_fn_grad)
+                                                         sample.output_process_fn_grad)
                 result = fn(*primals)
                 cotangents = tree_map(lambda x: torch.randn_like(x), result)
 
@@ -882,12 +891,12 @@ def test_vjpvmap(self, device, dtype, op):
 
                 self.assertEqual(result_vjps, expected_vjps)
 
+
 class InplaceError(Exception):
     def __repr__(self):
         return "Decomposition Tensor with no elem was created (probably due to an in-place op)"
 
 
-
 def ref_vjp_no_create(f, *primals):
     result = f(*primals)
 
@@ -896,12 +905,18 @@ def wrapped(cotangents):
 
     return result, wrapped
 
+
 run_decompositions = set()
 run_ops = set()
+
+
 class TestDecompositionOpInfo(TestCase):
 
     @unittest.skipIf(IS_FBCODE, "__torch_dispatch__ is buggy")
-    @ops(functorch_lagging_op_db + additional_op_db, allowed_dtypes=[torch.float32, torch.float64, torch.float16, torch.bfloat16] + [*integral_types()] )
+    @ops(
+        functorch_lagging_op_db + additional_op_db,
+        allowed_dtypes=[torch.float32, torch.float64, torch.float16, torch.bfloat16] + [*integral_types()]
+    )
     # entries in here need don't work and need to be fixed.
     # Each one of these is a bug (or needs to be investigated)
     @skipOps('TestDecompositionOpInfo', 'test_decomposition', {
@@ -912,8 +927,8 @@ class TestDecompositionOpInfo(TestCase):
         xfail('to_sparse'),
         skip('tensor_split'),
         skip('mvlgamma'),
-        skip('tanh', device_type='cuda'), # cuda bfloat16 failure
-        skip('nn.functional.tanhshrink', device_type='cuda'), # cuda bfloat16 failure
+        skip('tanh', device_type='cuda'),  # cuda bfloat16 failure
+        skip('nn.functional.tanhshrink', device_type='cuda'),  # cuda bfloat16 failure
         skip('eig'),
         skip('nn.functional.dropout'),
         skip('_masked.softmin'),
@@ -928,27 +943,30 @@ class TestDecompositionOpInfo(TestCase):
     def test_decomposition(self, device, dtype, op):
         # copied from common_utils.py
         dtype_precisions = {
-            torch.float16    : (0.001, 1e-5),
-            torch.bfloat16   : (0.016, 1e-5),
-            torch.float32    : (1.3e-6, 1e-5),
-            torch.float64    : (1e-7, 1e-7),
-            torch.complex32  : (0.001, 1e-5),
-            torch.complex64  : (1.3e-6, 1e-5),
-            torch.complex128 : (1e-7, 1e-7),
+            torch.float16: (0.001, 1e-5),
+            torch.bfloat16: (0.016, 1e-5),
+            torch.float32: (1.3e-6, 1e-5),
+            torch.float64: (1e-7, 1e-7),
+            torch.complex32: (0.001, 1e-5),
+            torch.complex64: (1.3e-6, 1e-5),
+            torch.complex128: (1e-7, 1e-7),
         }
         # Returns the "default" rtol and atol for comparing scalars or
         # tensors of the given dtypes.
+
         def _getDefaultRtolAndAtol(dtype0, dtype1):
             rtol = max(dtype_precisions.get(dtype0, (0, 0))[0],
-                    dtype_precisions.get(dtype1, (0, 0))[0])
+                       dtype_precisions.get(dtype1, (0, 0))[0])
             atol = max(dtype_precisions.get(dtype0, (0, 0))[1],
-                    dtype_precisions.get(dtype1, (0, 0))[1])
+                       dtype_precisions.get(dtype1, (0, 0))[1])
             return rtol, atol
+
         def op_assert_equal(op, a, b):
             assert a.dtype == b.dtype
             # Some ops, like those involving reductions, are fundamentally non-decomposable with precision guarantees
             tol_table = {
-                (torch.bfloat16, aten._softmax_backward_data): (0.016, 1e-2), # aggghhhhhhhhhh I hate reductions and floating point
+                # aggghhhhhhhhhh I hate reductions and floating point
+                (torch.bfloat16, aten._softmax_backward_data): (0.016, 1e-2),
                 (torch.bfloat16, aten._log_softmax_backward_data): (0.016, 1e-2),
                 # (torch.float16, aten.im2col_backward): (0.016, 1e-2),
             }
@@ -960,8 +978,10 @@ def op_assert_equal(op, a, b):
             assert torch.allclose(a, b, rtol=rtol, atol=atol), msg
 
         # We check the correctness of each decomposition right after running it.
-        # So, when we encounter a decomposition, we run the function normally, and then run the decomposition, and ensure they're identical.
-        # The way this is implemented, there could .... technically be an exponential blow up, but it's probably fine for now.
+        # So, when we encounter a decomposition, we run the function normally, and
+        # then run the decomposition, and ensure they're identical.
+        # The way this is implemented, there could .... technically be an exponential blow up,
+        # but it's probably fine for now.
         class DecompositionTensor(torch.Tensor):
             elem: torch.Tensor
 
@@ -986,6 +1006,7 @@ def __repr__(self):
             def __torch_dispatch__(cls, func, types, args=(), kwargs=None):
                 global run_ops
                 run_ops.add(func)
+
                 def unwrap_tensor(e):
                     if isinstance(e, DecompositionTensor):
                         if not hasattr(e, 'elem'):
@@ -993,10 +1014,8 @@ def unwrap_tensor(e):
                         return e.elem
                     return e
 
-
                 real_out = func(*tree_map(unwrap_tensor, args), **tree_map(unwrap_tensor, kwargs))
 
-
                 if func in decomposition_table and func != torch.ops.aten.detach:
                     upcast_table = set([
                         aten.tanh_backward,
@@ -1012,10 +1031,10 @@ def upcast_tensor(x):
                             x = x.to(dtype=torch.float32)
                         return x
                     # Theoretically, most PyTorch ops compute intermediates as fp32. But this breaks some ops...
-                    if func in upcast_table: 
-                        args = tree_map(upcast_tensor, args) 
+                    if func in upcast_table:
+                        args = tree_map(upcast_tensor, args)
                         kwargs = tree_map(upcast_tensor, kwargs)
-                    decomp_out =  decomposition(*args, **kwargs)
+                    decomp_out = decomposition(*args, **kwargs)
                     real_out_flat = tree_flatten(real_out)[0]
                     decomp_out_flat = tree_flatten(decomp_out)[0]
                     assert(len(real_out_flat) == len(decomp_out_flat))
@@ -1030,7 +1049,7 @@ def wrap_tensor(e):
                     if e is None:
                         return DecompositionTensor(torch.empty(()))
                     return DecompositionTensor(e) if type(e) == torch.Tensor else e
-                wrapped_out =  tree_map(wrap_tensor, real_out)
+                wrapped_out = tree_map(wrap_tensor, real_out)
                 return wrapped_out
 
         if dtype not in op.supported_dtypes(dtype):
@@ -1054,23 +1073,22 @@ def wrap_tensor(x):
             for sample_input in samples:
                 if _requires_grad:
                     fn, primals = normalize_op_input_output(func, sample_input)
-                    primals = tree_map (lambda x: x.abs() if isinstance(x, torch.Tensor) else x, primals)
+                    primals = tree_map(lambda x: x.abs() if isinstance(x, torch.Tensor) else x, primals)
 
                     decomp_out, decomp_vjp_fn = ref_vjp_no_create(fn, *tree_map(wrap_tensor, primals))
                     cotangents = tree_map(lambda x: torch.randn_like(x), decomp_out)
 
-                    decomp_grads = decomp_vjp_fn(cotangents)
+                    _ = decomp_vjp_fn(cotangents)
 
                 else:
                     args = [sample_input.input] + list(sample_input.args)
                     kwargs = sample_input.kwargs
-                    orig_out = func(*args, **kwargs)
+                    _ = func(*args, **kwargs)
 
                     args = tree_map(wrap_tensor, args)
                     kwargs = tree_map(wrap_tensor, kwargs)
                     decomp_out = func(*args, **kwargs)
 
-
         except InplaceError:
             self.skipTest("op is inplace")
             return
@@ -1087,14 +1105,15 @@ def wrap_tensor(x):
     def test_placeholder(self):
         global run_ops, run_decompositions
         with open('op_analysis/run_ops.txt', 'w') as f:
-            def get_names(l):
-                return sorted([x.__name__ for x in l])
+            def get_names(inpt):
+                return sorted([x.__name__ for x in inpt])
             for op in get_names(run_ops):
                 f.write(f'{op}\n')
         with open('op_analysis/run_decompositions.txt', 'w') as f:
             for op in get_names(run_decompositions):
                 f.write(f'{op}\n')
 
+
 only_for = ("cpu", "cuda")
 instantiate_device_type_tests(TestOperators, globals(), only_for=only_for)
 instantiate_device_type_tests(TestDecompositionOpInfo, globals(), only_for=only_for)
diff --git a/test/test_pythonkey.py b/test/test_pythonkey.py
index 27e5ba257..2ca76113c 100644
--- a/test/test_pythonkey.py
+++ b/test/test_pythonkey.py
@@ -7,36 +7,23 @@
 from torch.testing._internal.common_utils import TestCase, run_tests
 import torch
 import torch.nn as nn
-import torch.nn.functional as F
 import torch.utils._pytree as pytree
 import unittest
-import functools
-import itertools
 import warnings
-import math
-from typing import Callable, Type
-from torch.testing._internal.common_device_type import instantiate_device_type_tests, \
-    skipCUDAIfNoMagma, onlyCPU
-import types
-from functools import partial, wraps
-import copy
-
-import functorch
+from torch.testing._internal.common_device_type import instantiate_device_type_tests
 from functorch import (
-    grad, vjp, vmap, jacrev, grad_and_value,
+    grad, vjp, vmap, jacrev,
     make_fx
 )
 from functorch.compile import (
     nnc_jit, compiled_function, compiled_module,
-    partition_with_recompute_fwd_in_bwd, pythonkey_decompose, decomposition_table, aot_function, aot_module
+    partition_with_recompute_fwd_in_bwd, pythonkey_decompose, aot_function, aot_module
 )
 
-from torch.testing._internal.common_device_type import ops, onlyCPU
+from torch.testing._internal.common_device_type import ops
 from functorch_lagging_op_db import functorch_lagging_op_db
 from functorch_additional_op_db import additional_op_db
 from common_utils import (
-    get_fallback_and_vmap_exhaustive,
-    opinfo_in_dict,
     xfail,
     skip,
     skipOps,
@@ -53,7 +40,7 @@
                   UserWarning)
 
 # NB: numpy is a testing dependency!
-import numpy as np
+
 
 class TestPythonKey(TestCase):
     def test_make_fx(self, device):
@@ -191,6 +178,7 @@ def f(x):
     @unittest.skipIf(not USE_TORCHVISION, "test requires torchvision")
     def test_resnet18_backward_trace(self, device):
         mod = torchvision.models.resnet18()
+
         def f(x):
             out = mod(x)
             out.sum().backward()
@@ -216,17 +204,19 @@ def f(x):
     xfail('linalg.eigvals'),
     xfail('nn.functional.ctc_loss'),
     xfail('nn.functional.fractional_max_pool3d', device_type='cpu'),
-    xfail('randn_like'), # randomness
-    xfail('rand_like'), # randomness
-    xfail('randint_like'), # randomness
-    skip('new_empty'), # nondeterministic
-    skip('empty_like'), # nondeterministic
-    skip('linalg.lstsq', 'grad_oriented'), # flaky
+    xfail('randn_like'),  # randomness
+    xfail('rand_like'),  # randomness
+    xfail('randint_like'),  # randomness
+    skip('new_empty'),  # nondeterministic
+    skip('empty_like'),  # nondeterministic
+    skip('linalg.lstsq', 'grad_oriented'),  # flaky
 }
+
+
 class TestPythonKeyOperatorsOpInfo(TestCase):
     @ops(functorch_lagging_op_db + additional_op_db, allowed_dtypes=(torch.float,))
     @skipOps('TestPythonKeyOperatorsOpInfo', 'test_make_fx_exhaustive', make_fx_failures
-    )
+             )
     def test_make_fx_exhaustive(self, device, dtype, op):
 
         def f(args, kwargs):
@@ -247,15 +237,17 @@ def f(args, kwargs):
                     arg.uniform_(0, 1)
             try:
                 old_out = f(args, kwargs)
-            except:
+            except Exception:
                 continue
             new_out = new_f(args, kwargs)
             self.assertEqual(new_out, old_out)
             pass
 
+
 def _nop_compile(x, _):
     return x
 
+
 def _outs_and_grads(fn, inps):
     outs = fn(*inps)
     [out.sum().backward(retain_graph=True) for out in pytree.tree_flatten(outs)[0]]
@@ -265,7 +257,6 @@ def _outs_and_grads(fn, inps):
     return outs, grads
 
 
-
 class TestAOTAutograd(TestCase):
     def verify_aot_autograd(self, f, inp):
         if isinstance(f, nn.Module):
@@ -295,12 +286,6 @@ def f(a, b):
         inp = [torch.randn(3, 3, requires_grad=True), torch.randn(3, 3)]
         self.verify_aot_autograd(f, inp)
 
-    def test_multi_output_list(self):
-        def f(a, b):
-            return [a + b, a - b]
-        inp = [torch.randn(3, 3, requires_grad=True), torch.randn(3, 3)]
-        self.verify_aot_autograd(f, inp)
-
     def test_output_dict(self):
         def f(x):
             return {'a': x, 'b': x}
@@ -311,7 +296,7 @@ def f(x, y):
             return {'a': x, 'b': y + x}
         inp = [torch.randn(3, requires_grad=True), torch.randn(3)]
         self.verify_aot_autograd(f, inp)
-        
+
         def f(x):
             new_d = {}
             for k in x:
@@ -320,7 +305,6 @@ def f(x):
         inp = [{'a': torch.randn(3, requires_grad=True), 'b': torch.randn(3, requires_grad=True)}]
         self.verify_aot_autograd(f, inp)
 
-
     def test_module(self):
         mod = nn.Sequential(nn.Linear(32, 32), nn.ReLU())
         compiled_mod = compiled_module(mod, _nop_compile, _nop_compile)
@@ -338,6 +322,7 @@ def test_batchnorm(self):
         x = torch.ones(1, 4, 2, 2)
         mod(x).sum().backward()
 
+
 class TestEagerFusionOpInfo(TestCase):
     @ops(functorch_lagging_op_db + additional_op_db, allowed_dtypes=(torch.float,))
     # entries in here need don't work and need to be fixed.
@@ -372,7 +357,6 @@ def f(args, kwargs):
         if not op.supports_autograd:
             return
         sample_inputs_itr = op.sample_inputs(device, dtype, requires_grad=True)
-        new_f = None
         for sample_input in sample_inputs_itr:
             args = [sample_input.input] + list(sample_input.args)
             kwargs = sample_input.kwargs
@@ -394,7 +378,7 @@ def f(x):
             def get_grads(args):
                 return pytree.tree_map(lambda x: x.grad, args)
 
-            compiled_f = compiled_function(f, lambda x,_: x, lambda x,_: x)
+            compiled_f = compiled_function(f, lambda x, _: x, lambda x, _: x)
 
             reset_grads()
             compiled_f(args, kwargs).sum().backward()
@@ -419,6 +403,7 @@ def create_new_arg(x):
             orig_grad = get_grads(args)
             self.assertEqual(orig_grad, compiled_grad)
 
+
 class TestPartitioning(TestCase):
     def test_recompute_partitioning(self):
         def fn(a, b):
@@ -433,7 +418,10 @@ def fn(a, b):
         # Compiled function calculation
         res_a = ref_a.clone().detach().requires_grad_(True)
         res_b = ref_b.clone().detach().requires_grad_(True)
-        compile_fn = lambda x, _ : x
+
+        def compile_fn(x, _):
+            return x
+
         compiled_fn = compiled_function(fn, compile_fn, compile_fn, partition_with_recompute_fwd_in_bwd)
         res = compiled_fn(res_a, res_b)
         res.sum().backward()
diff --git a/test/test_vmap.py b/test/test_vmap.py
index 2b9a040d3..bfef2b980 100644
--- a/test/test_vmap.py
+++ b/test/test_vmap.py
@@ -52,7 +52,7 @@ def __exit__(self, *ignored):
 class TestVmapAPI(TestCase):
     def test_non_tensor_output_raises(self):
         with self.assertRaisesRegex(ValueError, "got type <class 'float'> as a return"):
-            output = vmap(lambda x: 3.14)(torch.ones(3))
+            vmap(lambda x: 3.14)(torch.ones(3))
 
         def multiple_outputs(x):
             return x, 3
@@ -471,7 +471,6 @@ def f(x):
             vmap(f)(torch.randn(3))
 
     def test_accepts_nested_inputs(self):
-        B0 = 2
         x = torch.randn(2, 3)
         y = torch.randn(2, 3)
 
@@ -541,7 +540,6 @@ def bar(x, yz):
             return x * yz[0] * yz[1]
 
         x = torch.randn(2, 3)
-        y = torch.randn(2, 3)
 
         # the following are errors in jax (and will always be errors)
         msg = 'Got in_dim=0 for an input but the input is of type'
@@ -584,7 +582,7 @@ def test_fallback_does_not_warn_by_default(self):
         x = torch.randn(11)
         y = torch.randn(11)
         with warnings.catch_warnings(record=True) as wa:
-            result = vmap(op)(x, y)
+            vmap(op)(x, y)
             # The single warning here is the "vmap is experimental"
             # warning, not a warning from the vmap fallback path.
             self.assertEqual(len(wa), 1)
@@ -599,17 +597,17 @@ def test_fallback_warns_when_warnings_are_enabled(self):
         y = torch.randn(11)
         with warnings.catch_warnings(record=True) as wa:
             with EnableVmapFallbackWarnings():
-                result = vmap(op)(x, y)
+                vmap(op)(x, y)
             self.assertEqual(len(wa), 2)
             self.assertRegex(str(wa[-1].message), FALLBACK_REGEX)
 
     def _assert_uses_vmap_fallback(self, vmap_args, inputs):
         return
-        with warnings.catch_warnings(record=True) as wa:
-            with EnableVmapFallbackWarnings():
-                result = vmap(*vmap_args)(*inputs)
-            self.assertEqual(len(wa), 2)
-            self.assertRegex(str(wa[-1].message), FALLBACK_REGEX)
+        # with warnings.catch_warnings(record=True) as wa:
+        #     with EnableVmapFallbackWarnings():
+        #         result = vmap(*vmap_args)(*inputs)
+        #     self.assertEqual(len(wa), 2)
+        #     self.assertRegex(str(wa[-1].message), FALLBACK_REGEX)
 
     def test_fallback_zero_dim(self):
         # NB: One day we will implement a batching rule for torch.atan2.
@@ -790,7 +788,7 @@ def test_inplace_fallback_nary_same_levels(self):
         x_orig = torch.randn(B0, B1, B2, 5)
         x = x_orig.clone()
         y = torch.randn(B0, B1, B2)
-        result = vmap(vmap(vmap(op)))(x, y)
+        vmap(vmap(vmap(op)))(x, y)
         self.assertEqual(x, outplace_op(x_orig, y.view(B0, B1, B2, 1)))
 
     # ("Fallback isInplaceVmapCompatible check is broken")
@@ -801,7 +799,7 @@ def test_inplace_fallback_nary_different_levels(self):
         # path on another operator to avoid bitrot.
         op = Tensor.atan2_
         outplace_op = torch.atan2
-        B0, B1, B2 = 2, 3, 5
+        B0, B1 = 2, 3
 
         x = torch.rand(B0, 7)
         y = torch.rand(7)
@@ -1115,6 +1113,8 @@ def randp1(size, device='cpu', dtype=torch.float):
 #
 # check_view: Test if the first returned output is a view of the first input
 # check_propagates_grad: Test if the operation propagates gradients.
+
+
 def _vmap_test(self, op, inputs, in_dims=0, out_dims=0,
                check_view=False, check_propagates_grad=True):
     result = vmap(op, in_dims, out_dims)(*inputs)
@@ -1143,9 +1143,11 @@ def _vmap_test(self, op, inputs, in_dims=0, out_dims=0,
     result_as_tuple = (result,) if op_has_single_return else result
     self.assertTrue(result[0].requires_grad)
 
+
 def should_allow_vmap_fallback_usage(fn):
     return getattr(fn, '_allow_vmap_fallback_usage', False)
 
+
 def allowVmapFallbackUsage(fn):
     fn._allow_vmap_fallback_usage = True
     return fn
@@ -1160,6 +1162,8 @@ def allowVmapFallbackUsage(fn):
 #
 # NB: TestVmapBase is a nested class. This prevents test runners from picking
 # it up and running it.
+
+
 class Namespace:
     class TestVmapBase(TestCase):
         def __init__(self, method_name='runTest'):
@@ -1174,18 +1178,18 @@ def __init__(self, method_name='runTest'):
                         self._wrap_method_with_vmap_fallback_check(test_method))
 
         def _wrap_method_with_vmap_fallback_check(self, method):
-            msg = (
-                'Expected the test to not invoke the vmap fallback path, i.e., '
-                'all of the operators being tested in this test should have batching '
-                'rules implemented. If you are intentionally testing something to '
-                'do with the fallback path, use allowVmapFallbackUsage. Otherwise, '
-                'please make sure that batching rules are implemented for the '
-                'operator(s) being tested.'
-            )
+            # msg = (
+            #     'Expected the test to not invoke the vmap fallback path, i.e., '
+            #     'all of the operators being tested in this test should have batching '
+            #     'rules implemented. If you are intentionally testing something to '
+            #     'do with the fallback path, use allowVmapFallbackUsage. Otherwise, '
+            #     'please make sure that batching rules are implemented for the '
+            #     'operator(s) being tested.'
+            # )
 
             @functools.wraps(method)
             def wrapper(self, *args, **kwargs):
-                with warnings.catch_warnings(record=True) as wa:
+                with warnings.catch_warnings(record=True):
                     warnings.simplefilter('always')
                     with EnableVmapFallbackWarnings():
                         method(*args, **kwargs)
@@ -1221,6 +1225,7 @@ def uses_fallback(self):
             with self.assertRaises(AssertionError):
                 uses_fallback(self)
 
+
 def _make_case(op, input_getter=TensorFactory.randn):
     return (op, input_getter)
 
@@ -1310,7 +1315,6 @@ def clone_contiguous(x):
         self.assertTrue(y.is_contiguous())
         self.assertTrue(y[0][0].is_contiguous())
 
-
         msg = r'only supported with memory_format torch.preserve_format or torch.contiguous_format'
         with self.assertRaisesRegex(RuntimeError, msg):
             vmap(lambda x: x.clone(memory_format=torch.channels_last))(torch.randn(B0))
@@ -2032,7 +2036,7 @@ def baz(x, memory_format):
     def test_unsqueeze(self):
         op = torch.unsqueeze
         test = self._vmap_view_test
-        B0, B1, B2 = 7, 11, 13
+        B0, B1 = 7, 11
 
         # unsqueeze dim 0
         test(op, (torch.rand(B0, 2, 5), 0), in_dims=(0, None))
@@ -2260,7 +2264,6 @@ def op(*tensors):
         test(vmap(get_op(0), in_dims=(0, 0)),
              (torch.rand(B1, 2), torch.rand(B0, B1, 2)), in_dims=(None, 0))
 
-
     def test_slice(self):
         test = self._vmap_view_test
         B0, B1, B2 = 7, 11, 13
@@ -2698,12 +2701,13 @@ def test_conv2d(self):
             for loop_out, batched_out in get_fallback_and_vmap_exhaustive(conv_fn, arg_values, kwarg_values):
                 self.assertEqual(loop_out, batched_out)
 
-            mod2 = torch.nn.Conv2d(4, 8, kernel_size=3, groups=2, stride=3, padding=1, dilation = 2)
-            arg_values = [torch.randn(inp_shape), mod.weight, mod.bias]
+            mod2 = conv_mod(4, 8, kernel_size=3, groups=2, stride=3, padding=1, dilation=2)
+            arg_values = [torch.randn(inp_shape), mod2.weight, mod2.bias]
+            kwarg_values = dict(groups=2, stride=3, padding=1, dilation=2)
             for loop_out, batched_out in get_fallback_and_vmap_exhaustive(conv_fn, arg_values, kwarg_values):
                 self.assertEqual(loop_out, batched_out)
 
-            arg_values = [torch.randn(inp_shape), mod.weight, None]
+            arg_values = [torch.randn(inp_shape), mod2.weight, None]
             for loop_out, batched_out in get_fallback_and_vmap_exhaustive(conv_fn, arg_values, kwarg_values):
                 self.assertEqual(loop_out, batched_out)
 
@@ -2718,6 +2722,7 @@ def test_one_hot(self):
 
     def test_conj_bit(self):
         x = torch.tensor([1+1j, 2+1j])
+
         def foo(x):
             assert not x.is_conj()
             y = x.conj()
@@ -2758,6 +2763,7 @@ def test_foobar_parametrize(self, op):
     def test_parametrize_multiple(self, op1, op2):
         pass
 
+
 instantiate_parametrized_tests(TestVmapOperators)
 
 
@@ -2765,6 +2771,7 @@ def construct_v(output, batch_size):
     return torch.randn(batch_size, *output.shape,
                        dtype=output.dtype, device=output.device)
 
+
 def as_tuple(x):
     if isinstance(x, tuple):
         return x
@@ -2773,10 +2780,12 @@ def as_tuple(x):
     else:
         return x,
 
+
 def differentiable(args):
     return tuple(arg for arg in as_tuple(args)
                  if isinstance(arg, torch.Tensor) and arg.requires_grad)
 
+
 def _get_rand_no_zeros(*args, **kwargs):
     requires_grad = kwargs.get('requires_grad', False)
     kwargs_without_requires_grad = kwargs.copy()
@@ -2784,6 +2793,7 @@ def _get_rand_no_zeros(*args, **kwargs):
     result = torch.rand(*args, **kwargs_without_requires_grad)
     return result.clamp_min_(0.1).requires_grad_(requires_grad)
 
+
 class TestVmapBatchedGradient(Namespace.TestVmapBase):
     def _vmap_test(self, *args, **kwargs):
         return _vmap_test(self, *args, **kwargs)
@@ -3048,7 +3058,6 @@ def test_diagonal(self, device):
         x = torch.randn(3, 4, 5, device=device, requires_grad=True)
         self._batched_grad_test(lambda x: x.diagonal(0, -1, -2), (x,))
 
-
     @allowVmapFallbackUsage
     def test_unrelated_output(self, device):
         B0 = 3
@@ -3078,6 +3087,7 @@ def vjp(v):
         result = vmap(vjp)(gy)
         self.assertEqual(result, torch.zeros(B0, *x.shape, device=device))
 
+
 class TestVmapOperatorsOpInfo(TestCase):
     vmap_fail = {
         # These are ops that we can't generate fallbacks for
@@ -3133,6 +3143,7 @@ class TestVmapOperatorsOpInfo(TestCase):
         xfail('nn.functional.glu'),
         xfail('cartesian_prod'),
     }
+
     @ops(functorch_lagging_op_db + additional_op_db, allowed_dtypes=(torch.float,))
     @skipOps('TestVmapOperatorsOpInfo', 'test_vmap_exhaustive', vmap_fail)
     def test_vmap_exhaustive(self, device, dtype, op):
@@ -3321,7 +3332,8 @@ def test_isinf(self, device):
         test(self, op, (x,), in_dims=(0))
 
     def test_foo_like(self, device):
-        test = functools.partial(_vmap_test, check_propagates_grad=False)
+        # vfdev-5: Probably, we can remove this line. Flake8 reported as unused
+        # test = functools.partial(_vmap_test, check_propagates_grad=False)
 
         B, N, C, H, W = 2, 3, 24, 5, 7
         for op in [torch.ones_like, torch.zeros_like, torch.randn_like, torch.rand_like]:
@@ -3355,20 +3367,23 @@ def test_group_norm(self, device):
         test(self, op, (x, 4, weight, bias), in_dims=(0, None, 0, 0))
 
     def test_index_put(self, device):
-        test = functools.partial(_vmap_test, check_propagates_grad=False)
+        # vfdev-5: Probably, we can remove this line. Flake8 reported as unused
+        # test = functools.partial(_vmap_test, check_propagates_grad=False)
+
+        x = torch.arange(3 * 4 * 5).reshape(3, 4, 5)
 
-        x = torch.arange(3*4*5).reshape(3,4,5)
         def f(x, y, z):
             x[y] = z
             return x
+
         x = torch.randn(3, 4, 5)
         y = torch.zeros((3, 2)).long()
         z = torch.randn(3, 2, 5)
         base = f(x[0], y[0], z[0])
-        self.assertEqual(vmap(f, in_dims=(0,0,0))(x, y, z)[0], base)
-        self.assertEqual(vmap(f, in_dims=(0,None,None))(x, y[0], z[0])[0], base)
-        self.assertEqual(vmap(f, in_dims=(0,None,0))(x, y[0], z)[0], base)
-        self.assertEqual(vmap(f, in_dims=(0,0,None))(x, y, z[0])[0], base)
+        self.assertEqual(vmap(f, in_dims=(0, 0, 0))(x, y, z)[0], base)
+        self.assertEqual(vmap(f, in_dims=(0, None, None))(x, y[0], z[0])[0], base)
+        self.assertEqual(vmap(f, in_dims=(0, None, 0))(x, y[0], z)[0], base)
+        self.assertEqual(vmap(f, in_dims=(0, 0, None))(x, y, z[0])[0], base)
 
     @parametrize('training', [True, False])
     @parametrize('track_running_stats', [True, False])
diff --git a/test/xfail_suggester.py b/test/xfail_suggester.py
index 12a66223a..3d8c82596 100644
--- a/test/xfail_suggester.py
+++ b/test/xfail_suggester.py
@@ -1,5 +1,4 @@
 import re
-import pprint
 
 """
 Instructions:
@@ -12,7 +11,8 @@
     lines = f.readlines()
 
 failed = [line for line in lines if line.startswith('FAILED')]
-p = re.compile('FAILED test/test_\w+.py::\w+::(\S+)')
+p = re.compile('FAILED test/test_\w+.py::\w+::(\S+)')  # noqa: W605
+
 
 def get_failed_test(line):
     m = p.match(line)
@@ -20,6 +20,7 @@ def get_failed_test(line):
         return None
     return m.group(1)
 
+
 base_names = {
     'test_grad_',
     'test_vjp_',
@@ -41,9 +42,11 @@ def get_failed_test(line):
 
 suggested_xfails = {}
 
+
 def remove_device_dtype(test):
     return '_'.join(test.split('_')[:-2])
 
+
 def belongs_to_base(test, base):
     if not test.startswith(base):
         return False
@@ -53,6 +56,7 @@ def belongs_to_base(test, base):
             return False
     return True
 
+
 def sanitize_base(base):
     if base.startswith('nn_functional_'):
         base = f'nn.functional.{base[len("nn_functional_"):]}'
@@ -64,12 +68,14 @@ def sanitize_base(base):
         base = f'_masked.{base[len("_masked_"):]}'
     return base
 
+
 def any_starts_with(strs, thing):
     for s in strs:
         if s.startswith(thing):
             return True
     return False
 
+
 def get_suggested_xfails(base, tests):
     result = []
     tests = [test[len(base):] for test in tests if
@@ -93,6 +99,7 @@ def get_suggested_xfails(base, tests):
         result.append(f"skip('{sanitized_base}'),")
     return result
 
+
 result = {base: get_suggested_xfails(base, failed_tests) for base in base_names}
 for k, v in result.items():
     print('=' * 50)