flake

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
 

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(

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:

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)

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 = [

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()