Fixes

Author: vivekmig

captum/attr/_core/deep_lift.py

@@ -582,7 +582,9 @@ def __init__(self, model: Module, multiply_by_inputs: bool = True) -> None:
     def attribute(
         self,
         inputs: TensorOrTupleOfTensorsGeneric,
-        baselines: Union[BaselineType, Callable[..., TensorOrTupleOfTensorsGeneric]],
+        baselines: Union[
+            TensorOrTupleOfTensorsGeneric, Callable[..., TensorOrTupleOfTensorsGeneric]
+        ],
         target: TargetType = None,
         additional_forward_args: Any = None,
         return_convergence_delta: Literal[False] = False,
@@ -593,7 +595,9 @@ def attribute(
     def attribute(
         self,
         inputs: TensorOrTupleOfTensorsGeneric,
-        baselines: Union[BaselineType, Callable[..., TensorOrTupleOfTensorsGeneric]],
+        baselines: Union[
+            TensorOrTupleOfTensorsGeneric, Callable[..., TensorOrTupleOfTensorsGeneric]
+        ],
         target: TargetType = None,
         additional_forward_args: Any = None,
         *,

tests/attr/test_deeplift_classification.py

@@ -155,6 +155,11 @@ def softmax_classification(
         target: TargetType,
     ) -> None:
         # TODO add test cases for multiple different layers
+        if isinstance(attr_method, DeepLiftShap):
+            assert isinstance(
+                baselines, Tensor
+            ), "Non-tensor baseline not supported for DeepLiftShap"
+
         model.zero_grad()
         attributions, delta = attr_method.attribute(
             input, baselines=baselines, target=target, return_convergence_delta=True

tests/helpers/basic.py

@@ -19,7 +19,9 @@ def copy_args(*args, **kwargs):
     return copy_args
 
 
-def assertTensorAlmostEqual(test, actual, expected, delta=0.0001, mode="sum"):
+def assertTensorAlmostEqual(
+    test, actual, expected, delta: float = 0.0001, mode: str = "sum"
+):
     assert isinstance(actual, torch.Tensor), (
         "Actual parameter given for " "comparison must be a tensor."
     )
@@ -57,7 +59,9 @@ def assertTensorAlmostEqual(test, actual, expected, delta=0.0001, mode="sum"):
         raise ValueError("Mode for assertion comparison must be one of `max` or `sum`.")
 
 
-def assertTensorTuplesAlmostEqual(test, actual, expected, delta=0.0001, mode="sum"):
+def assertTensorTuplesAlmostEqual(
+    test, actual, expected, delta: float = 0.0001, mode: str = "sum"
+):
     if isinstance(expected, tuple):
         assert len(actual) == len(
             expected

tests/robust/test_attack_comparator.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 import collections
-from typing import List
+from typing import Dict, List, Tuple, Union
 
 import torch
 from captum.robust import AttackComparator, FGSM
@@ -202,7 +202,9 @@ def test_attack_comparator_with_additional_args(self) -> None:
         attack_comp.reset()
         self.assertEqual(len(attack_comp.summary()), 0)
 
-    def _compare_results(self, obtained, expected) -> None:
+    def _compare_results(
+        self, obtained: Union[Dict, Tuple, Tensor], expected: Union[Dict, Tuple, Tensor]
+    ) -> None:
         if isinstance(expected, dict):
             self.assertIsInstance(obtained, dict)
             for key in expected: