Adjust indices in LayerAttributor mask for individual neurons (#1531)

Summary:

With `enable_cross_tensor_attribution=True` for `FeatureAblation`/`FeaturePermutation`, ids/indices in the masks are now "global"

Reviewed By: cyrjano

Differential Revision: D71778355

Author: sarahtranfb

captum/_utils/common.py

@@ -194,17 +194,18 @@ def _is_mask_valid(mask: Tensor, inp: Tensor) -> bool:
 def _format_feature_mask(
     feature_mask: Union[None, Tensor, Tuple[Tensor, ...]],
     inputs: Tuple[Tensor, ...],
+    start_idx: int = 0,
 ) -> Tuple[Tensor, ...]:
     """
     Format a feature mask into a tuple of tensors.
     The `inputs` should be correctly formatted first
     If `feature_mask` is None, assign each non-batch dimension with a consecutive
-    integer from 0.
+    integer from `start_idx`.
     If `feature_mask` is a tensor, wrap it in a tuple.
     """
     if feature_mask is None:
         formatted_mask = []
-        current_num_features = 0
+        current_num_features = start_idx
         for inp in inputs:
             # the following can handle empty tensor where numel is 0
             # empty tensor will be added to the feature mask

captum/testing/helpers/basic_models.py

@@ -662,6 +662,16 @@ def forward(self, x1: Tensor, x2: Tensor, x3: Tensor, scale: int):
         return self.model(scale * (x1 + x2 + x3))
 
 
+class BasicModel_MultiLayer_TupleInput(nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.model = BasicModel_MultiLayer()
+
+    @no_type_check
+    def forward(self, x: Tuple[Tensor, Tensor, Tensor]) -> Tensor:
+        return self.model(x[0] + x[1] + x[2])
+
+
 class BasicModel_MultiLayer_MultiInput_with_Future(nn.Module):
     def __init__(self) -> None:
         super().__init__()

tests/attr/neuron/test_neuron_ablation.py

@@ -83,8 +83,8 @@ def test_multi_input_ablation_with_mask(self) -> None:
         inp2 = torch.tensor([[20.0, 50.0, 30.0], [0.0, 100.0, 0.0]])
         inp3 = torch.tensor([[0.0, 100.0, 10.0], [2.0, 10.0, 3.0]])
         mask1 = torch.tensor([[1, 1, 1], [0, 1, 0]])
-        mask2 = torch.tensor([[0, 1, 2]])
-        mask3 = torch.tensor([[0, 1, 2], [0, 0, 0]])
+        mask2 = torch.tensor([[3, 4, 2]])
+        mask3 = torch.tensor([[5, 6, 7], [5, 5, 5]])
         expected = (
             [[492.0, 492.0, 492.0], [200.0, 200.0, 200.0]],
             [[80.0, 200.0, 120.0], [0.0, 400.0, 0.0]],