Move methods out of FeatureAblation and make them free methods. (#1663)

Summary:

This diff moves methods out of the `FeatureAblation` class and makes them free methods. The changes include creating a new method `_parse_forward_out` to force forward output type assertion and conversion, and modifying the `add_one_back` module to use the new methods. The `attr/fb/add_one_back.py` file has been modified to use the new methods. The `attr/fb/within_group_utils.py` file has also been modified to use the new methods. The `attr/fb/test_within_group_utils.py` file has been

Differential Revision: D86785624

Author: cyrjano

captum/attr/_core/feature_ablation.py

@@ -37,6 +37,91 @@
 logger: logging.Logger = logging.getLogger(__name__)
 
 
+def _parse_forward_out(forward_output: object) -> Tensor:
+    """
+    A temp wrapper for global _run_forward util to force forward output
+    type assertion & conversion.
+    Remove after the strict logic is supported by all attr classes
+    """
+    if isinstance(forward_output, Tensor):
+        return forward_output
+
+    output_type = type(forward_output)
+    assert output_type is int or output_type is float, (
+        "the return of forward_func must be a tensor, int, or float,"
+        f" received: {forward_output}"
+    )
+
+    # using python built-in type as torch dtype
+    # int -> torch.int64, float -> torch.float64
+    # ref: https://github.com/pytorch/pytorch/pull/21215
+    return torch.tensor(forward_output, dtype=cast(dtype, output_type))
+
+
+def process_initial_eval(
+    initial_eval: Tensor,
+    inputs: TensorOrTupleOfTensorsGeneric,
+    use_weights: bool = False,
+) -> Tuple[List[Tensor], List[Tensor], Tensor, Tensor, int, dtype]:
+
+    initial_eval = _parse_forward_out(initial_eval)
+
+    # number of elements in the output of forward_func
+    n_outputs = initial_eval.numel() if isinstance(initial_eval, Tensor) else 1
+
+    # flatten eval outputs into 1D (n_outputs)
+    # add the leading dim for n_feature_perturbed
+    flattened_initial_eval = initial_eval.reshape(1, -1)
+
+    # Initialize attribution totals and counts
+    attrib_type = flattened_initial_eval.dtype
+
+    total_attrib = [
+        # attribute w.r.t each output element
+        torch.zeros(
+            (n_outputs,) + input.shape[1:],
+            dtype=attrib_type,
+            device=input.device,
+        )
+        for input in inputs
+    ]
+
+    # Weights are used in cases where ablations may be overlapping.
+    weights = []
+    if use_weights:
+        weights = [
+            torch.zeros((n_outputs,) + input.shape[1:], device=input.device).float()
+            for input in inputs
+        ]
+
+    return (
+        total_attrib,
+        weights,
+        initial_eval,
+        flattened_initial_eval,
+        n_outputs,
+        attrib_type,
+    )
+
+
+def format_result(
+    total_attrib: List[Tensor],
+    weights: List[Tensor],
+    is_inputs_tuple: bool,
+    use_weights: bool,
+) -> Union[Tensor, Tuple[Tensor, ...]]:
+    """Normalizes attributions by weights if enabled and formats output as single tensor or tuple."""
+    # Divide total attributions by counts and return formatted attributions
+    if use_weights:
+        attrib = tuple(
+            single_attrib.float() / weight
+            for single_attrib, weight in zip(total_attrib, weights)
+        )
+    else:
+        attrib = tuple(total_attrib)
+    return _format_output(is_inputs_tuple, attrib)
+
+
 class FeatureAblation(PerturbationAttribution):
     r"""
     A perturbation based approach to computing attribution, involving
@@ -331,9 +416,8 @@ def attribute(
                 flattened_initial_eval,
                 n_outputs,
                 attrib_type,
-            ) = self._process_initial_eval(
-                initial_eval,
-                formatted_inputs,
+            ) = process_initial_eval(
+                initial_eval, formatted_inputs, use_weights=self.use_weights
             )
 
             total_attrib, weights = self._attribute_with_cross_tensor_feature_masks(
@@ -358,7 +442,9 @@ def attribute(
 
         return cast(
             TensorOrTupleOfTensorsGeneric,
-            self._generate_result(total_attrib, weights, is_inputs_tuple),
+            format_result(
+                total_attrib, weights, is_inputs_tuple, use_weights=self.use_weights
+            ),
         )
 
     def _attribute_with_cross_tensor_feature_masks(
@@ -586,8 +672,8 @@ def _initial_eval_to_processed_initial_eval_fut(
                     "initial_eval_to_processed_initial_eval_fut: "
                     "initial_eval should be a Tensor"
                 )
-            result = self._process_initial_eval(
-                initial_eval_processed, formatted_inputs
+            result = process_initial_eval(
+                initial_eval_processed, formatted_inputs, use_weights=self.use_weights
             )
 
         except FeatureAblationFutureError as e:
@@ -886,10 +972,8 @@ def _generate_async_result_cross_tensor(
             )
 
         result_fut = collect_all(accumulate_fut_list).then(
-            lambda x: self._generate_result(
-                total_attrib,
-                weights,
-                is_inputs_tuple,
+            lambda x: format_result(
+                total_attrib, weights, is_inputs_tuple, use_weights=self.use_weights
             )
         )
 
@@ -955,70 +1039,6 @@ def _eval_fut_to_ablated_out_fut_cross_tensor(
             ) from e
         return total_attrib, weights
 
-    def _parse_forward_out(self, forward_output: Tensor) -> Tensor:
-        """
-        A temp wrapper for global _run_forward util to force forward output
-        type assertion & conversion.
-        Remove after the strict logic is supported by all attr classes
-        """
-        if isinstance(forward_output, Tensor):
-            return forward_output
-
-        output_type = type(forward_output)
-        assert output_type is int or output_type is float, (
-            "the return of forward_func must be a tensor, int, or float,"
-            f" received: {forward_output}"
-        )
-
-        # using python built-in type as torch dtype
-        # int -> torch.int64, float -> torch.float64
-        # ref: https://github.com/pytorch/pytorch/pull/21215
-        return torch.tensor(forward_output, dtype=cast(dtype, output_type))
-
-    def _process_initial_eval(
-        self,
-        initial_eval: Tensor,
-        inputs: TensorOrTupleOfTensorsGeneric,
-    ) -> Tuple[List[Tensor], List[Tensor], Tensor, Tensor, int, dtype]:
-        initial_eval = self._parse_forward_out(initial_eval)
-
-        # number of elements in the output of forward_func
-        n_outputs = initial_eval.numel() if isinstance(initial_eval, Tensor) else 1
-
-        # flatten eval outputs into 1D (n_outputs)
-        # add the leading dim for n_feature_perturbed
-        flattened_initial_eval = initial_eval.reshape(1, -1)
-
-        # Initialize attribution totals and counts
-        attrib_type = flattened_initial_eval.dtype
-
-        total_attrib = [
-            # attribute w.r.t each output element
-            torch.zeros(
-                (n_outputs,) + input.shape[1:],
-                dtype=attrib_type,
-                device=input.device,
-            )
-            for input in inputs
-        ]
-
-        # Weights are used in cases where ablations may be overlapping.
-        weights = []
-        if self.use_weights:
-            weights = [
-                torch.zeros((n_outputs,) + input.shape[1:], device=input.device).float()
-                for input in inputs
-            ]
-
-        return (
-            total_attrib,
-            weights,
-            initial_eval,
-            flattened_initial_eval,
-            n_outputs,
-            attrib_type,
-        )
-
     def _process_ablated_out_full(
         self,
         modified_eval: Tensor,
@@ -1033,7 +1053,7 @@ def _process_ablated_out_full(
         attrib_type: dtype,
         perturbations_per_eval: int,
     ) -> Tuple[List[Tensor], List[Tensor]]:
-        modified_eval = self._parse_forward_out(modified_eval)
+        modified_eval = _parse_forward_out(modified_eval)
         # if perturbations_per_eval > 1, the output shape must grow with
         # input and not be aggregated
         current_batch_size = inputs[0].shape[0]
@@ -1086,19 +1106,3 @@ def _process_ablated_out_full(
             total_attrib[i] += (eval_diff * mask.to(attrib_type)).sum(dim=0)
 
         return total_attrib, weights
-
-    def _generate_result(
-        self,
-        total_attrib: List[Tensor],
-        weights: List[Tensor],
-        is_inputs_tuple: bool,
-    ) -> Union[Tensor, Tuple[Tensor, ...]]:
-        # Divide total attributions by counts and return formatted attributions
-        if self.use_weights:
-            attrib = tuple(
-                single_attrib.float() / weight
-                for single_attrib, weight in zip(total_attrib, weights)
-            )
-        else:
-            attrib = tuple(total_attrib)
-        return _format_output(is_inputs_tuple, attrib)

captum/testing/helpers/basic_models.py

@@ -581,7 +581,6 @@ def forward(
             self.relu(lin1_out)
         else:
             relu_out = self.relu(lin1_out)
-        # pyre-fixme [29]: `typing.Type[Future]` is not a function
         result = Future()
         lin2_out = self.linear2(relu_out)
         if multidim_output:

tests/attr/test_feature_ablation.py

@@ -12,7 +12,11 @@
 import torch
 from captum._utils.common import _construct_future_forward
 from captum._utils.typing import BaselineType, TargetType, TensorOrTupleOfTensorsGeneric
-from captum.attr._core.feature_ablation import FeatureAblation
+from captum.attr._core.feature_ablation import (
+    _parse_forward_out,
+    FeatureAblation,
+    format_result,
+)
 from captum.attr._core.noise_tunnel import NoiseTunnel
 from captum.attr._utils.attribution import Attribution
 from captum.testing.helpers import BaseTest
@@ -595,7 +599,6 @@ def slow_set_future(fut: torch.futures.Future[Tensor], value: Tensor) -> None:
             fut.set_result(out)
 
         def forward_func(inp: Tensor) -> torch.futures.Future[Tensor]:
-            # pyre-fixme[29]: `typing.Type[torch.futures.Future]` is not a function.
             fut: torch.futures.Future[Tensor] = torch.futures.Future()
             t = threading.Thread(target=slow_set_future, args=(fut, inp))
             t.start()
@@ -900,5 +903,115 @@ def _ablation_test_assert(
                 assertTensorAlmostEqual(self, attributions, expected_ablation)
 
 
+class TestParseForwardOutput(BaseTest):
+
+    def test_parse_forward_out_tensor_passthrough(self) -> None:
+        input_tensor = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
+        result = _parse_forward_out(input_tensor)
+
+        self.assertIs(result, input_tensor)
+        assertTensorAlmostEqual(self, result, input_tensor)
+
+    def test_parse_forward_out_python_int(self) -> None:
+        input_value = 42
+        result = _parse_forward_out(input_value)
+
+        self.assertIsInstance(result, Tensor)
+        self.assertEqual(result.dtype, torch.int64)
+        assertTensorAlmostEqual(self, result, torch.tensor(42))
+
+    def test_parse_forward_out_python_float(self) -> None:
+        input_value = 3.14
+        result = _parse_forward_out(input_value)
+
+        self.assertIsInstance(result, Tensor)
+        self.assertEqual(result.dtype, torch.float64)
+        assertTensorAlmostEqual(self, result, torch.tensor(3.14))
+
+    def test_parse_forward_out_invalid_none(self) -> None:
+        with self.assertRaises(AssertionError):
+            _parse_forward_out(None)
+
+
+class TestFormatResult(BaseTest):
+
+    def test_format_result_single_tensor_no_weights(self) -> None:
+        total_attrib = [torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])]
+        weights = []
+        is_inputs_tuple = False
+        use_weights = False
+
+        result = format_result(total_attrib, weights, is_inputs_tuple, use_weights)
+
+        self.assertIsInstance(result, Tensor)
+        assert isinstance(result, Tensor)  # Type narrowing for pyre
+        self.assertEqual(result.shape, (2, 3))
+        assertTensorAlmostEqual(
+            self, result, torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
+        )
+
+    def test_format_result_tuple_output_no_weights(self) -> None:
+        total_attrib = [
+            torch.tensor([[1.0, 2.0], [3.0, 4.0]]),
+            torch.tensor([[5.0, 6.0], [7.0, 8.0]]),
+        ]
+        weights = []
+        is_inputs_tuple = True
+        use_weights = False
+
+        result = format_result(total_attrib, weights, is_inputs_tuple, use_weights)
+
+        self.assertIsInstance(result, tuple)
+        self.assertEqual(len(result), 2)
+        assertTensorAlmostEqual(self, result[0], torch.tensor([[1.0, 2.0], [3.0, 4.0]]))
+        assertTensorAlmostEqual(self, result[1], torch.tensor([[5.0, 6.0], [7.0, 8.0]]))
+
+    def test_format_result_single_tensor_with_weights(self) -> None:
+        total_attrib = [torch.tensor([[10.0, 20.0, 30.0], [40.0, 50.0, 60.0]])]
+        weights = [torch.tensor([[2.0, 4.0, 5.0], [8.0, 10.0, 12.0]])]
+        is_inputs_tuple = False
+        use_weights = True
+
+        result = format_result(total_attrib, weights, is_inputs_tuple, use_weights)
+
+        self.assertIsInstance(result, Tensor)
+        expected = torch.tensor([[5.0, 5.0, 6.0], [5.0, 5.0, 5.0]])
+        assertTensorAlmostEqual(self, result, expected)
+
+    def test_format_result_tuple_output_with_weights(self) -> None:
+        total_attrib = [
+            torch.tensor([[10.0, 20.0], [30.0, 40.0]]),
+            torch.tensor([[50.0, 60.0], [70.0, 80.0]]),
+        ]
+        weights = [
+            torch.tensor([[2.0, 4.0], [5.0, 8.0]]),
+            torch.tensor([[10.0, 12.0], [14.0, 16.0]]),
+        ]
+        is_inputs_tuple = True
+        use_weights = True
+
+        result = format_result(total_attrib, weights, is_inputs_tuple, use_weights)
+
+        self.assertIsInstance(result, tuple)
+        self.assertEqual(len(result), 2)
+        assertTensorAlmostEqual(self, result[0], torch.tensor([[5.0, 5.0], [6.0, 5.0]]))
+        assertTensorAlmostEqual(self, result[1], torch.tensor([[5.0, 5.0], [5.0, 5.0]]))
+
+    def test_format_result_integer_dtype_no_weights(self) -> None:
+        total_attrib = [torch.tensor([[1, 2, 3], [4, 5, 6]], dtype=torch.int32)]
+        weights = []
+        is_inputs_tuple = False
+        use_weights = False
+
+        result = format_result(total_attrib, weights, is_inputs_tuple, use_weights)
+
+        self.assertIsInstance(result, Tensor)
+        assert isinstance(result, Tensor)  # Type narrowing for pyre
+        self.assertEqual(result.dtype, torch.int32)
+        assertTensorAlmostEqual(
+            self, result, torch.tensor([[1, 2, 3], [4, 5, 6]], dtype=torch.int32)
+        )
+
+
 if __name__ == "__main__":
     unittest.main()