From 23f3100df9adf8df872432fd92a5acf10f44aff8 Mon Sep 17 00:00:00 2001
From: Thomas George <tfjgeorge@gmail.com>
Date: Fri, 24 Nov 2023 23:03:33 +0100
Subject: [PATCH] flake

---
 nngeometry/generator/jacobian/grads.py      | 19 +++++++----
 nngeometry/generator/jacobian/grads_conv.py |  7 ++--
 nngeometry/layercollection.py               |  5 ++-
 tests/tasks.py                              |  2 +-
 tests/test_jacobian.py                      | 38 ++++++++++++++-------
 tests/test_jacobian_kfac.py                 |  2 +-
 6 files changed, 48 insertions(+), 25 deletions(-)

diff --git a/nngeometry/generator/jacobian/grads.py b/nngeometry/generator/jacobian/grads.py
index fe2bf76..77414b8 100644
--- a/nngeometry/generator/jacobian/grads.py
+++ b/nngeometry/generator/jacobian/grads.py
@@ -1,12 +1,19 @@
 import torch
 import torch.nn.functional as F
 
-from nngeometry.layercollection import (Affine1dLayer, BatchNorm1dLayer,
-                                        BatchNorm2dLayer, Conv2dLayer,
-                                        ConvTranspose2dLayer, Cosine1dLayer,
-                                        GroupNormLayer, LinearLayer,
-                                        WeightNorm1dLayer, WeightNorm2dLayer,
-                                        Conv1dLayer)
+from nngeometry.layercollection import (
+    Affine1dLayer,
+    BatchNorm1dLayer,
+    BatchNorm2dLayer,
+    Conv2dLayer,
+    ConvTranspose2dLayer,
+    Cosine1dLayer,
+    GroupNormLayer,
+    LinearLayer,
+    WeightNorm1dLayer,
+    WeightNorm2dLayer,
+    Conv1dLayer,
+)
 
 from .grads_conv import conv2d_backward, convtranspose2d_backward, conv1d_backward
 
diff --git a/nngeometry/generator/jacobian/grads_conv.py b/nngeometry/generator/jacobian/grads_conv.py
index c7b025c..565ea5b 100644
--- a/nngeometry/generator/jacobian/grads_conv.py
+++ b/nngeometry/generator/jacobian/grads_conv.py
@@ -128,8 +128,11 @@ def conv1d_backward_using_unfold(mod, x, gy):
     gy_s = gy.size()
     bs = gy_s[0]
     x_unfold = F.unfold(
-        x.unsqueeze(2), kernel_size=ks, stride=(1, mod.stride[0]),
-        padding=(0, mod.padding[0]), dilation=(1, mod.dilation[0])
+        x.unsqueeze(2),
+        kernel_size=ks,
+        stride=(1, mod.stride[0]),
+        padding=(0, mod.padding[0]),
+        dilation=(1, mod.dilation[0]),
     )
     x_unfold_s = x_unfold.size()
     return torch.bmm(
diff --git a/nngeometry/layercollection.py b/nngeometry/layercollection.py
index 344174b..e3924de 100644
--- a/nngeometry/layercollection.py
+++ b/nngeometry/layercollection.py
@@ -238,14 +238,13 @@ def __eq__(self, other):
             and self.kernel_size == other.kernel_size
         )
 
+
 class Conv1dLayer(AbstractLayer):
     def __init__(self, in_channels, out_channels, kernel_size, bias=True):
         self.in_channels = in_channels
         self.out_channels = out_channels
         self.kernel_size = kernel_size
-        self.weight = Parameter(
-            out_channels, in_channels, kernel_size[0]
-        )
+        self.weight = Parameter(out_channels, in_channels, kernel_size[0])
         if bias:
             self.bias = Parameter(out_channels)
         else:
diff --git a/tests/tasks.py b/tests/tasks.py
index bfa98d4..32ea77c 100644
--- a/tests/tasks.py
+++ b/tests/tasks.py
@@ -487,7 +487,7 @@ def output_fn(input, target):
 class Conv1dNet(nn.Module):
     def __init__(self, normalization="none"):
         super(Conv1dNet, self).__init__()
-        if normalization != 'none':
+        if normalization != "none":
             raise NotImplementedError
         self.normalization = normalization
         self.conv1 = nn.Conv1d(1, 6, 3, 3)
diff --git a/tests/test_jacobian.py b/tests/test_jacobian.py
index db6b234..c188364 100644
--- a/tests/test_jacobian.py
+++ b/tests/test_jacobian.py
@@ -1,21 +1,35 @@
 import pytest
 import torch
-from tasks import (get_batchnorm_conv_linear_task,
-                   get_batchnorm_fc_linear_task, get_conv_gn_task,
-                   get_conv_skip_task, get_conv_task,
-                   get_fullyconnect_affine_task, get_fullyconnect_cosine_task,
-                   get_fullyconnect_onlylast_task, get_fullyconnect_task,
-                   get_fullyconnect_wn_task, get_linear_conv_task,
-                   get_linear_fc_task, get_small_conv_transpose_task,
-                   get_small_conv_wn_task, get_conv1d_task)
+from tasks import (
+    get_batchnorm_conv_linear_task,
+    get_batchnorm_fc_linear_task,
+    get_conv_gn_task,
+    get_conv_skip_task,
+    get_conv_task,
+    get_fullyconnect_affine_task,
+    get_fullyconnect_cosine_task,
+    get_fullyconnect_onlylast_task,
+    get_fullyconnect_task,
+    get_fullyconnect_wn_task,
+    get_linear_conv_task,
+    get_linear_fc_task,
+    get_small_conv_transpose_task,
+    get_small_conv_wn_task,
+    get_conv1d_task,
+)
 from utils import check_ratio, check_tensors
 
 from nngeometry.generator import Jacobian
 from nngeometry.object.fspace import FMatDense
-from nngeometry.object.map import (PullBackDense, PushForwardDense,
-                                   PushForwardImplicit)
-from nngeometry.object.pspace import (PMatBlockDiag, PMatDense, PMatDiag,
-                                      PMatImplicit, PMatLowRank, PMatQuasiDiag)
+from nngeometry.object.map import PullBackDense, PushForwardDense, PushForwardImplicit
+from nngeometry.object.pspace import (
+    PMatBlockDiag,
+    PMatDense,
+    PMatDiag,
+    PMatImplicit,
+    PMatLowRank,
+    PMatQuasiDiag,
+)
 from nngeometry.object.vector import PVector, random_fvector, random_pvector
 
 linear_tasks = [
diff --git a/tests/test_jacobian_kfac.py b/tests/test_jacobian_kfac.py
index aa528f9..2188197 100644
--- a/tests/test_jacobian_kfac.py
+++ b/tests/test_jacobian_kfac.py
@@ -145,7 +145,7 @@ def test_jacobian_kfac_vs_pblockdiag():
     """
     for get_task, mult in zip(
         [get_conv1dnet_kfc_task, get_convnet_kfc_task, get_fullyconnect_kfac_task],
-        [3., 15.0, 1.0]
+        [3.0, 15.0, 1.0],
     ):
         loader, lc, parameters, model, function, n_output = get_task()