Fixed python code formatting and added flake8 setup (#346)

* Fixed python code formatting and added flake8 setup

* Fixes config.yaml

* Added missing setup.cfg

* Removed flake8 job from circle ci and setup GHA

* More flake8 fixes

* Fixed test_conv2d

* Fixed failing flake8

Author: vfdev-5

.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

codegen/codegen_outofplacebatching.py

@@ -6,48 +6,55 @@
 
 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
     if typ == 'const Tensor &':
         return True
     return False
 
+
 def is_optional_tensor(typ: str) -> bool:
     if typ == 'c10::optional<Tensor>':
         return True
     if typ == 'const c10::optional<Tensor> &':
         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,13 +75,15 @@ 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;
     optional<int64_t> {name}_bdim;
     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',

codegen/gen_plumbing.py

@@ -1,5 +1,4 @@
 import argparse
-import re
 
 from codegen_outofplacebatching import deindent, get_signatures, gen_unwraps
 

examples/compilation/eager_fusion.py

@@ -1,15 +1,17 @@
 from functorch import compiled_function, tvm_compile
-import torch.nn as nn
 import torch
 from functools import partial
 import time
 import torch.utils
 
 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()

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)

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,23 +42,28 @@ 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)
     loss.backward()
     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)
 

examples/compilation/simple_function.py

@@ -8,16 +8,19 @@
 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)
 grad_nnc = nnc_jit(grad_pt, skip_specialization=True)
 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))