flake8

nngeometry/generator/jacobian/grads_conv.py

@@ -83,11 +83,6 @@ def conv_backward(
     return weight_bgrad
 
 
-def conv1d_backward(*args, **kwargs):
-    """Computes per-example gradients for nn.Conv1d layers."""
-    return conv_backward(*args, nd=1, **kwargs)
-
-
 def conv2d_backward_using_conv(mod, x, gy):
     """Computes per-example gradients for nn.Conv2d layers."""
     return conv_backward(

nngeometry/generator/jacobian/jacobian.py

@@ -41,7 +41,10 @@ def __init__(
         self.centering = centering
 
         if function is None:
-            function = lambda *x: model(x[0])
+
+            def function(*x):
+                return model(x[0])
+
         self.function = function
 
         if layer_collection is None:

nngeometry/object/__init__.py

@@ -1,7 +1,15 @@
 from .fspace import FMatDense
 from .map import PullBackDense, PushForwardDense, PushForwardImplicit
-from .pspace import (PMatBlockDiag, PMatDense, PMatDiag, PMatEKFAC,
-                     PMatImplicit, PMatKFAC, PMatLowRank, PMatQuasiDiag)
+from .pspace import (
+    PMatBlockDiag,
+    PMatDense,
+    PMatDiag,
+    PMatEKFAC,
+    PMatImplicit,
+    PMatKFAC,
+    PMatLowRank,
+    PMatQuasiDiag,
+)
 from .vector import FVector, PVector
 
 __all__ = [

tests/tasks.py

@@ -370,21 +370,6 @@ def get_fullyconnect_affine_task():
     return get_fullyconnect_task(normalization="affine")
 
 
-def get_conv_task(normalization="none"):
-    train_set = get_mnist()
-    train_set = Subset(train_set, range(70))
-    train_loader = DataLoader(dataset=train_set, batch_size=30, shuffle=False)
-    net = ConvNet(normalization=normalization)
-    to_device_model(net)
-    net.eval()
-
-    def output_fn(input, target):
-        return net(to_device(input))
-
-    layer_collection = LayerCollection.from_model(net)
-    return (train_loader, layer_collection, net.parameters(), net, output_fn, 3)
-
-
 def get_conv_bn_task():
     return get_conv_task(normalization="batch_norm")
 

tests/test_jacobian.py

@@ -833,6 +833,7 @@ def test_bn_eval_mode():
         model.train()
         with pytest.raises(RuntimeError):
             FMat_dense = FMatDense(generator=generator, examples=loader)
+            FMat_dense.get_dense_tensor()
 
 
 def test_example_passing():

tests/test_jacobian_kfac.py

@@ -3,7 +3,6 @@
 import pytest
 import torch
 import torch.nn as nn
-import torch.nn.functional as F
 from tasks import (
     get_conv_task,
     get_fullyconnect_task,
@@ -196,8 +195,6 @@ def test_jacobian_kfac():
         # Test mv
         mv_direct = torch.mv(G_kfac_split, random_v.get_flat_representation())
         mv_kfac = M_kfac.mv(random_v)
-        print(mv_direct.size(), lc.layers)
-        pvec = PVector(layer_collection=lc, vector_repr=mv_direct)
         check_tensors(mv_direct, mv_kfac.get_flat_representation())
 
         # Test vTMv

tests/test_metrics.py

@@ -1,9 +1,14 @@
 import pytest
 import torch
 import torch.nn.functional as tF
-from tasks import (device, get_conv_gn_task, get_conv_task,
-                   get_fullyconnect_segm_task, get_fullyconnect_task,
-                   to_device)
+from tasks import (
+    device,
+    get_conv_gn_task,
+    get_conv_task,
+    get_fullyconnect_segm_task,
+    get_fullyconnect_task,
+    to_device,
+)
 from test_jacobian import get_output_vector, update_model
 
 from nngeometry.metrics import FIM, FIM_MonteCarlo

tests/test_pickle.py

@@ -4,8 +4,13 @@
 from utils import check_tensors
 
 from nngeometry.generator import Jacobian
-from nngeometry.object.pspace import (PMatBlockDiag, PMatDense, PMatDiag,
-                                      PMatLowRank, PMatQuasiDiag)
+from nngeometry.object.pspace import (
+    PMatBlockDiag,
+    PMatDense,
+    PMatDiag,
+    PMatLowRank,
+    PMatQuasiDiag,
+)
 from nngeometry.object.vector import PVector
 
 

tests/test_vector.py

@@ -5,8 +5,7 @@
 from utils import check_ratio, check_tensors
 
 from nngeometry.layercollection import LayerCollection
-from nngeometry.object.vector import (PVector, random_pvector,
-                                      random_pvector_dict)
+from nngeometry.object.vector import PVector, random_pvector, random_pvector_dict
 
 
 class ConvNet(nn.Module):