[ADD] hessian_is_zero in SlicingDerivatives

Added an additional test case with small enough input to run the Hessian property test.
f-dangel · Jan 4, 2022 · 48e4bbe · 48e4bbe
1 parent 2235f0c
commit 48e4bbe
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 1 deletion.
diff --git a/backpack/core/derivatives/slicing.py b/backpack/core/derivatives/slicing.py
@@ -42,3 +42,6 @@ def no_slice_batch_axis(module: Slicing):
 
         if slice_batch_axis != slice(None):
             raise ValueError("Slicing the batch axis is not supported.")
+
+    def hessian_is_zero(self, module: Slicing) -> bool:  # noqa: D102
+        return True
diff --git a/test/core/derivatives/derivatives_test.py b/test/core/derivatives/derivatives_test.py
@@ -408,7 +408,11 @@ def test_ea_jac_t_mat_jac_prod(problem: DerivativesTestProblem, request) -> None
 
 
 @fixture(
-    params=PROBLEMS + BATCH_NORM_PROBLEMS + RNN_PROBLEMS + EMBEDDING_PROBLEMS,
+    params=PROBLEMS
+    + BATCH_NORM_PROBLEMS
+    + RNN_PROBLEMS
+    + EMBEDDING_PROBLEMS
+    + CUSTOM_SLICING_MODULE_PROBLEMS,
     ids=lambda p: p.make_id(),
 )
 def problem(request) -> DerivativesTestProblem:

diff --git a/test/core/derivatives/slicing_settings.py b/test/core/derivatives/slicing_settings.py
@@ -6,6 +6,10 @@
 from backpack.custom_module.slicing import Slicing
 
 CUSTOM_SLICING_SETTINGS = [
+    {
+        "module_fn": lambda: Slicing((slice(None), 0)),
+        "input_fn": lambda: rand(size=(2, 4, 2, 5)),
+    },
     {
         "module_fn": lambda: Slicing((slice(None),)),
         "input_fn": lambda: rand(size=(3, 4, 2, 5)),