Repository: MachineLearningLifeScience/stochman (language: python)

Issue:
nnh.leakyrelu seems incorrect
The Jacobian of leakyrelu seems to be the Jacobian of a relu, which isn't quite right.


Discussion / hints:

Fix patch:
diff --git a/stochman/nnj.py b/stochman/nnj.py
index 19a36bf..60b0294 100644
--- a/stochman/nnj.py
+++ b/stochman/nnj.py
@@ -301,8 +301,8 @@ def _jacobian(self, x: Tensor, val: Tensor) -> Tensor:
 
 class LeakyReLU(AbstractActivationJacobian, nn.LeakyReLU):
     def _jacobian(self, x: Tensor, val: Tensor) -> Tensor:
-        jac = torch.zeros_like(val)
-        jac[val.abs() < 1.0] = 1.0
+        jac = torch.ones_like(val)
+        jac[val < 0.0] = self.negative_slope
         return jac
 
 


Test patch:
diff --git a/tests/test_nnj.py b/tests/test_nnj.py
index 676c91e..92750e6 100644
--- a/tests/test_nnj.py
+++ b/tests/test_nnj.py
@@ -9,7 +9,7 @@
 
 _batch_size = 2
 _features = 5
-_dims = 10
+_dims = 6
 
 _linear_input = torch.randn(_batch_size, _features)
 _1d_conv_input = torch.randn(_batch_size, _features, _dims)
@@ -49,10 +49,10 @@ def _compare_jacobian(f: Callable, x: torch.Tensor) -> torch.Tensor:
                 nnj.ReLU(),
                 nnj.Sqrt(),
                 nnj.Hardshrink(),
-                nnj.LeakyReLU(),
             ),
             _linear_input,
         ),
+        (nnj.Sequential(nnj.Linear(_features, 2), nnj.LeakyReLU()), _linear_input),
         (nnj.Sequential(nnj.Linear(_features, 2), nnj.OneMinusX()), _linear_input),
         (nnj.Sequential(nnj.Conv1d(_features, 2, 5), nnj.ConvTranspose1d(2, _features, 5)), _1d_conv_input),
         (nnj.Sequential(nnj.Conv2d(_features, 2, 5), nnj.ConvTranspose2d(2, _features, 5)), _2d_conv_input),
@@ -88,7 +88,7 @@ def _compare_jacobian(f: Callable, x: torch.Tensor) -> torch.Tensor:
             nnj.Sequential(
                 nnj.Conv1d(_features, 2, 3),
                 nnj.Flatten(),
-                nnj.Linear(8 * 2, 5),
+                nnj.Linear(4 * 2, 5),
                 nnj.ReLU(),
             ),
             _1d_conv_input,
@@ -97,7 +97,7 @@ def _compare_jacobian(f: Callable, x: torch.Tensor) -> torch.Tensor:
             nnj.Sequential(
                 nnj.Conv2d(_features, 2, 3),
                 nnj.Flatten(),
-                nnj.Linear(8 * 8 * 2, 5),
+                nnj.Linear(4 * 4 * 2, 5),
                 nnj.ReLU(),
             ),
             _2d_conv_input,
@@ -106,7 +106,7 @@ def _compare_jacobian(f: Callable, x: torch.Tensor) -> torch.Tensor:
             nnj.Sequential(
                 nnj.Conv3d(_features, 2, 3),
                 nnj.Flatten(),
-                nnj.Linear(8 * 8 * 8 * 2, 5),
+                nnj.Linear(4 * 4 * 4 * 2, 5),
                 nnj.ReLU(),
             ),
             _3d_conv_input,