Fixes for Boolean Input Tensors

captum/_utils/common.py

@@ -106,7 +106,7 @@ def _zeros(inputs: Tuple[Tensor, ...]) -> Tuple[int, ...]:
     Takes a tuple of tensors as input and returns a tuple that has the same
     length as `inputs` with each element as the integer 0.
     """
-    return tuple(0 for input in inputs)
+    return tuple(0 if input.dtype is not torch.bool else False for input in inputs)
 
 
 def _format_baseline(

captum/attr/_core/feature_ablation.py

@@ -554,7 +554,7 @@ def _get_feature_range_and_mask(self, input, input_mask, **kwargs):
         )
 
     def _get_feature_counts(self, inputs, feature_mask, **kwargs):
-        """ return the numbers of input features """
+        """return the numbers of input features"""
         if not feature_mask:
             return tuple(inp[0].numel() if inp.numel() else 0 for inp in inputs)
 

captum/attr/_core/lime.py

@@ -571,15 +571,18 @@ def default_from_interp_rep_transform(curr_sample, original_inputs, **kwargs):
     ), "Must provide baselines to use default interpretable representation transfrom"
     feature_mask = kwargs["feature_mask"]
     if isinstance(feature_mask, Tensor):
-        binary_mask = curr_sample[0][feature_mask].to(original_inputs.dtype)
-        return binary_mask * original_inputs + (1 - binary_mask) * kwargs["baselines"]
+        binary_mask = curr_sample[0][feature_mask].bool()
+        return (
+            binary_mask.to(original_inputs.dtype) * original_inputs
+            + (~binary_mask).to(original_inputs.dtype) * kwargs["baselines"]
+        )
     else:
         binary_mask = tuple(
-            curr_sample[0][feature_mask[j]] for j in range(len(feature_mask))
+            curr_sample[0][feature_mask[j]].bool() for j in range(len(feature_mask))
         )
         return tuple(
             binary_mask[j].to(original_inputs[j].dtype) * original_inputs[j]
-            + (1 - binary_mask[j].to(original_inputs[j].dtype)) * kwargs["baselines"][j]
+            + (~binary_mask[j]).to(original_inputs[j].dtype) * kwargs["baselines"][j]
             for j in range(len(feature_mask))
         )
 

captum/attr/_core/occlusion.py

@@ -375,5 +375,5 @@ def _get_feature_range_and_mask(
         return 0, feature_max, None
 
     def _get_feature_counts(self, inputs, feature_mask, **kwargs):
-        """ return the numbers of possible input features """
+        """return the numbers of possible input features"""
         return tuple(np.prod(counts).astype(int) for counts in kwargs["shift_counts"])

captum/attr/_core/shapley_value.py

@@ -424,8 +424,7 @@ def _perturbation_generator(
         target_repeated = _expand_target(target, perturbations_per_eval)
         for i in range(len(feature_permutation)):
             current_tensors = tuple(
-                current
-                * (torch.tensor(1) - (mask == feature_permutation[i]).to(current.dtype))
+                current * (~(mask == feature_permutation[i])).to(current.dtype)
                 + input * (mask == feature_permutation[i]).to(input.dtype)
                 for input, current, mask in zip(inputs, current_tensors, input_masks)
             )
@@ -478,7 +477,7 @@ def _perturbation_generator(
             )
 
     def _get_n_evaluations(self, total_features, n_samples, perturbations_per_eval):
-        """ return the total number of forward evaluations needed """
+        """return the total number of forward evaluations needed"""
         return math.ceil(total_features / perturbations_per_eval) * n_samples
 
 
@@ -740,7 +739,7 @@ def attribute(
         )
 
     def _get_n_evaluations(self, total_features, n_samples, perturbations_per_eval):
-        """ return the total number of forward evaluations needed """
+        """return the total number of forward evaluations needed"""
         return math.ceil(total_features / perturbations_per_eval) * math.factorial(
             total_features
         )

captum/metrics/_core/infidelity.py

@@ -69,7 +69,7 @@ def sub_infidelity_perturb_func_decorator(pertub_func: Callable) -> Callable:
         def default_perturb_func(
             inputs: TensorOrTupleOfTensorsGeneric, baselines: BaselineType = None
         ):
-            r""""""
+            r""" """
             inputs_perturbed = (
                 pertub_func(inputs, baselines)
                 if baselines is not None
@@ -398,7 +398,7 @@ def _generate_perturbations(
         """
 
         def call_perturb_func():
-            r""""""
+            r""" """
             baselines_pert = None
             inputs_pert: Union[Tensor, Tuple[Tensor, ...]]
             if len(inputs_expanded) == 1:

tests/attr/test_feature_ablation.py

@@ -18,6 +18,7 @@
     BasicModel_ConvNet_One_Conv,
     BasicModel_MultiLayer,
     BasicModel_MultiLayer_MultiInput,
+    BasicModelBoolInput,
     BasicModelWithSparseInputs,
 )
 
@@ -113,6 +114,29 @@ def test_simple_ablation_with_baselines(self) -> None:
             perturbations_per_eval=(1, 2, 3),
         )
 
+    def test_simple_ablation_boolean(self) -> None:
+        ablation_algo = FeatureAblation(BasicModelBoolInput())
+        inp = torch.tensor([[True, False, True]])
+        self._ablation_test_assert(
+            ablation_algo,
+            inp,
+            [[40.0, 40.0, 40.0]],
+            feature_mask=torch.tensor([[0, 0, 1]]),
+            perturbations_per_eval=(1, 2, 3),
+        )
+
+    def test_simple_ablation_boolean_with_baselines(self) -> None:
+        ablation_algo = FeatureAblation(BasicModelBoolInput())
+        inp = torch.tensor([[True, False, True]])
+        self._ablation_test_assert(
+            ablation_algo,
+            inp,
+            [[-40.0, -40.0, 0.0]],
+            feature_mask=torch.tensor([[0, 0, 1]]),
+            baselines=True,
+            perturbations_per_eval=(1, 2, 3),
+        )
+
     def test_multi_sample_ablation(self) -> None:
         ablation_algo = FeatureAblation(BasicModel_MultiLayer())
         inp = torch.tensor([[2.0, 10.0, 3.0], [20.0, 50.0, 30.0]], requires_grad=True)

tests/attr/test_lime.py

@@ -25,6 +25,7 @@
 from tests.helpers.basic_models import (
     BasicModel_MultiLayer,
     BasicModel_MultiLayer_MultiInput,
+    BasicModelBoolInput,
 )
 
 
@@ -146,6 +147,31 @@ def test_simple_lime_with_baselines(self) -> None:
             test_generator=True,
         )
 
+    def test_simple_lime_boolean(self) -> None:
+        net = BasicModelBoolInput()
+        inp = torch.tensor([[True, False, True]])
+        self._lime_test_assert(
+            net,
+            inp,
+            [31.42, 31.42, 30.90],
+            feature_mask=torch.tensor([[0, 0, 1]]),
+            perturbations_per_eval=(1, 2, 3),
+            test_generator=True,
+        )
+
+    def test_simple_lime_boolean_with_baselines(self) -> None:
+        net = BasicModelBoolInput()
+        inp = torch.tensor([[True, False, True]])
+        self._lime_test_assert(
+            net,
+            inp,
+            [-36.0, -36.0, 0.0],
+            feature_mask=torch.tensor([[0, 0, 1]]),
+            baselines=True,
+            perturbations_per_eval=(1, 2, 3),
+            test_generator=True,
+        )
+
     @unittest.mock.patch("sys.stderr", new_callable=io.StringIO)
     def test_simple_lime_with_show_progress(self, mock_stderr) -> None:
         net = BasicModel_MultiLayer()

tests/attr/test_shapley.py

@@ -13,6 +13,7 @@
 from tests.helpers.basic_models import (
     BasicModel_MultiLayer,
     BasicModel_MultiLayer_MultiInput,
+    BasicModelBoolInput,
 )
 
 
@@ -39,6 +40,29 @@ def test_simple_shapley_sampling_with_mask(self) -> None:
             perturbations_per_eval=(1, 2, 3),
         )
 
+    def test_simple_shapley_sampling_boolean(self) -> None:
+        net = BasicModelBoolInput()
+        inp = torch.tensor([[True, False, True]])
+        self._shapley_test_assert(
+            net,
+            inp,
+            [35.0, 35.0, 35.0],
+            feature_mask=torch.tensor([[0, 0, 1]]),
+            perturbations_per_eval=(1, 2, 3),
+        )
+
+    def test_simple_shapley_sampling_boolean_with_baseline(self) -> None:
+        net = BasicModelBoolInput()
+        inp = torch.tensor([[True, False, True]])
+        self._shapley_test_assert(
+            net,
+            inp,
+            [-40.0, -40.0, 0.0],
+            feature_mask=torch.tensor([[0, 0, 1]]),
+            baselines=True,
+            perturbations_per_eval=(1, 2, 3),
+        )
+
     def test_simple_shapley_sampling_with_baselines(self) -> None:
         net = BasicModel_MultiLayer()
         inp = torch.tensor([[20.0, 50.0, 30.0]])

tests/helpers/basic_models.py

@@ -362,6 +362,21 @@ def forward(
             return lin2_out
 
 
+class BasicModelBoolInput(nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.mod = BasicModel_MultiLayer()
+
+    def forward(
+        self,
+        x: Tensor,
+        add_input: Optional[Tensor] = None,
+        mult: float = 10.0,
+    ):
+        assert x.dtype is torch.bool, "Input must be boolean"
+        return self.mod(x.float() * mult, add_input)
+
+
 class BasicModel_MultiLayer_MultiInput(nn.Module):
     def __init__(self) -> None:
         super().__init__()