add PReLU

Author: SkafteNicki

stochman/nnj.py

@@ -320,6 +320,13 @@ def _jacobian(self, x: Tensor, val: Tensor) -> Tensor:
         return jac
 
 
+class PReLU(AbstractActivationJacobian, nn.PReLU):
+    def _jacobian(self, x: Tensor, val: Tensor) -> Tensor:
+        jac = torch.ones_like(val)
+        jac[x < 0.0] = self.weight
+        return jac
+
+
 class LeakyReLU(AbstractActivationJacobian, nn.LeakyReLU):
     def _jacobian(self, x: Tensor, val: Tensor) -> Tensor:
         jac = torch.ones_like(val)

tests/test_nnj.py

@@ -53,6 +53,7 @@ def _compare_jacobian(f: Callable, x: torch.Tensor) -> torch.Tensor:
         (nnj.Sequential(nnj.Linear(_features, 2), nnj.LeakyReLU()), _linear_input_shape),
         (nnj.Sequential(nnj.Linear(_features, 2), nnj.Tanh()), _linear_input_shape),
         (nnj.Sequential(nnj.Linear(_features, 2), nnj.OneMinusX()), _linear_input_shape),
+        (nnj.Sequential(nnj.Linear(_features, 2), nnj.PReLU()), _linear_input_shape),
         (
             nnj.Sequential(nnj.Conv1d(_features, 2, 5), nnj.ConvTranspose1d(2, _features, 5)),
             _1d_conv_input_shape,