Add test for influence utils (#1265)

Summary:
Pull Request resolved: #1265

As titled.
- Initial the test for influence utils.

Reviewed By: yucu

Differential Revision: D55375620

fbshipit-source-id: f9f9aee988b76c0ca64089c01b56aecb66e8f824

Author: cicichen01

captum/influence/_utils/common.py

@@ -108,7 +108,7 @@ def _jacobian_loss_wrt_inputs(
     batch).
 
     Args:
-        loss_fn (torch.nn.Module, Callable, or None): The loss function. If a library
+        loss_fn (torch.nn.Module, Callable): The loss function. If a library
                 defined loss function is provided, it would be expected to be a
                 torch.nn.Module. If a custom loss is provided, it can be either type,
                 but must behave as a library loss function would if `reduction='sum'`
@@ -131,24 +131,21 @@ def _jacobian_loss_wrt_inputs(
                 in the batch represented by `out`. This is a 2D tensor, where the
                 first dimension is the batch dimension.
     """
-    # TODO: allow loss_fn to be Callable
-    if isinstance(loss_fn, Module) and hasattr(loss_fn, "reduction"):
-        msg0 = "Please ensure that loss_fn.reduction is set to `sum` or `mean`"
-
-        assert loss_fn.reduction != "none", msg0
-        msg1 = (
-            f"loss_fn.reduction ({loss_fn.reduction}) does not match"
-            f"reduction type ({reduction_type}). Please ensure they are"
-            " matching."
-        )
-        assert loss_fn.reduction == reduction_type, msg1
-
     if reduction_type != "sum" and reduction_type != "mean":
         raise ValueError(
-            f"{reduction_type} is not a valid value for reduction_type. "
+            f"`{reduction_type}` is not a valid value for reduction_type. "
             "Must be either 'sum' or 'mean'."
         )
 
+    # TODO: allow loss_fn to be Callable
+    if isinstance(loss_fn, Module) and hasattr(loss_fn, "reduction"):
+        msg = (
+            f"loss_fn.reduction `{loss_fn.reduction}` does not match"
+            f"reduction type `{reduction_type}`. Please ensure they are"
+            " matching."
+        )
+        assert loss_fn.reduction == reduction_type, msg
+
     if _parse_version(torch.__version__) >= (1, 8, 0):
         input_jacobians = torch.autograd.functional.jacobian(
             lambda out: loss_fn(out, targets), out, vectorize=vectorize

tests/influence/_utils/__init__.py

@@ -0,0 +1,54 @@
+# (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary.
+
+# !/usr/bin/env python3
+
+import torch
+
+from captum.influence._utils.common import _jacobian_loss_wrt_inputs
+from tests.helpers import BaseTest
+from tests.helpers.basic import assertTensorAlmostEqual
+
+
+class TestCommon(BaseTest):
+    def setUp(self) -> None:
+        super().setUp()
+
+    def test_jacobian_loss_wrt_inputs(self) -> None:
+        with self.assertRaises(ValueError) as err:
+            _jacobian_loss_wrt_inputs(
+                torch.nn.BCELoss(reduction="sum"),
+                torch.tensor([-1.0, 1.0]),
+                torch.tensor([1.0]),
+                True,
+                "",
+            )
+        self.assertEqual(
+            "`` is not a valid value for reduction_type. "
+            "Must be either 'sum' or 'mean'.",
+            str(err.exception),
+        )
+
+        with self.assertRaises(AssertionError) as err:
+            _jacobian_loss_wrt_inputs(
+                torch.nn.BCELoss(reduction="sum"),
+                torch.tensor([-1.0, 1.0]),
+                torch.tensor([1.0]),
+                True,
+                "mean",
+            )
+        self.assertEqual(
+            "loss_fn.reduction `sum` does not matchreduction type `mean`."
+            " Please ensure they are matching.",
+            str(err.exception),
+        )
+
+        res = _jacobian_loss_wrt_inputs(
+            torch.nn.BCELoss(reduction="sum"),
+            torch.tensor([0.5, 1.0]),
+            torch.tensor([0.0, 1.0]),
+            True,
+            "sum",
+        )
+        assertTensorAlmostEqual(
+            self, res, torch.tensor([2.0, 0.0]), delta=0.0, mode="sum"
+        )