Fix total_forwards calculation in ablation/permutation for cross-tensor attribution

Summary: Previously, we could iterate over the feature masks and get the feature count from the ID range in the mask. Now mask ID/indices are global

Differential Revision: D72480911

Author: sarahtranfb

captum/attr/_core/feature_ablation.py

@@ -31,7 +31,10 @@
 from captum._utils.progress import progress, SimpleProgress
 from captum._utils.typing import BaselineType, TargetType, TensorOrTupleOfTensorsGeneric
 from captum.attr._utils.attribution import PerturbationAttribution
-from captum.attr._utils.common import _format_input_baseline
+from captum.attr._utils.common import (
+    _format_input_baseline,
+    get_total_features_from_mask,
+)
 from captum.log import log_usage
 from torch import dtype, Tensor
 from torch.futures import collect_all, Future
@@ -894,7 +897,9 @@ def _attribute_progress_setup(
             formatted_inputs, feature_mask, **kwargs
         )
         total_forwards = (
-            math.ceil(int(sum(feature_counts)) / perturbations_per_eval)
+            math.ceil(
+                get_total_features_from_mask(feature_mask) / perturbations_per_eval
+            )
             if enable_cross_tensor_attribution
             else sum(
                 math.ceil(count / perturbations_per_eval) for count in feature_counts

captum/attr/_utils/common.py

@@ -390,3 +390,16 @@ def _construct_default_feature_mask(
     total_features = current_num_features
     feature_mask = tuple(feature_mask)
     return feature_mask, total_features
+
+
+def get_total_features_from_mask(
+    feature_mask: Tuple[Tensor, ...],
+) -> int:
+    """
+    Return the numbers of input features based on the total unique
+    feature IDs/indices in the feature mask.
+    """
+    seen_idxs = set()
+    for mask in feature_mask:
+        seen_idxs |= set(torch.unique(mask).tolist())
+    return len(seen_idxs)

tests/utils/test_common.py

@@ -14,6 +14,7 @@
     parse_version,
     safe_div,
 )
+from captum.attr._utils.common import get_total_features_from_mask
 from captum.testing.helpers.basic import (
     assertTensorAlmostEqual,
     assertTensorTuplesAlmostEqual,
@@ -174,6 +175,16 @@ def test_get_max_feature_index(self) -> None:
 
         assert _get_max_feature_index(mask) == 100
 
+    def test_mask_unique_elem(self) -> None:
+        res = get_total_features_from_mask((torch.tensor([0, 0, 0]),))
+        self.assertEqual(res, 1)
+        res = get_total_features_from_mask((torch.tensor([0, 0, 4]),))
+        self.assertEqual(res, 2)
+        res = get_total_features_from_mask(
+            (torch.tensor([0, 0, 4]), torch.tensor([0, 4, 5]))
+        )
+        self.assertEqual(res, 3)
+
 
 class TestParseVersion(BaseTest):
     def test_parse_version_dev(self) -> None: