step for finite diff

tests/test_jacobian.py

@@ -1,22 +1,36 @@
 import pytest
 import torch
-from tasks import (get_batchnorm_conv_linear_task,
-                   get_batchnorm_fc_linear_task, get_conv1d_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)
+from tasks import (
+    get_batchnorm_conv_linear_task,
+    get_batchnorm_fc_linear_task,
+    get_conv1d_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,
+)
 from test_tasks.layernorm import get_layernorm_conv_task, get_layernorm_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 = [
@@ -100,7 +114,7 @@ def test_jacobian_pushforward_dense_nonlinear():
         )
         push_forward = PushForwardDense(generator=generator, examples=loader)
         dw = random_pvector(lc, device=device)
-        dw = 1e-4 / dw.norm() * dw
+        dw = 1e-5 / dw.norm() * dw
 
         doutput_lin = push_forward.mv(dw)
 
@@ -110,12 +124,6 @@ def test_jacobian_pushforward_dense_nonlinear():
 
         # This is non linear, so we don't expect the finite difference
         # estimate to be very accurate. We use a larger eps value
-        check_tensors(
-            output_after - output_before,
-            doutput_lin.get_flat_representation().t(),
-            eps=5e-3,
-            only_print_diff=True,
-        )
         check_tensors(
             output_after - output_before,
             doutput_lin.get_flat_representation().t(),