Fixes

Author: vivekmig

captum/_utils/common.py

@@ -3,7 +3,7 @@
 from enum import Enum
 from functools import reduce
 from inspect import signature
-from typing import Any, Callable, cast, Dict, List, overload, Tuple, Union
+from typing import Any, Callable, cast, Dict, List, overload, Sequence, Tuple, Union
 
 import numpy as np
 import torch
@@ -683,7 +683,7 @@ def _extract_device(
 
 
 def _reduce_list(
-    val_list: List[TupleOrTensorOrBoolGeneric],
+    val_list: Sequence[TupleOrTensorOrBoolGeneric],
     red_func: Callable[[List], Any] = torch.cat,
 ) -> TupleOrTensorOrBoolGeneric:
     """

captum/attr/_core/deep_lift.py

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

captum/influence/_core/tracincp_fast_rand_proj.py

@@ -867,7 +867,7 @@ class TracInCPFastRandProj(TracInCPFast):
     def __init__(
         self,
         model: Module,
-        final_fc_layer: Union[Module, str],
+        final_fc_layer: Module,
         train_dataset: Union[Dataset, DataLoader],
         checkpoints: Union[str, List[str], Iterator],
         checkpoints_load_func: Callable = _load_flexible_state_dict,
@@ -884,11 +884,9 @@ def __init__(
 
             model (torch.nn.Module): An instance of pytorch model. This model should
                     define all of its layers as attributes of the model.
-            final_fc_layer (torch.nn.Module or str): The last fully connected layer in
+            final_fc_layer (torch.nn.Module): The last fully connected layer in
                     the network for which gradients will be approximated via fast random
-                    projection method. Can be either the layer module itself, or the
-                    fully qualified name of the layer if it is a defined attribute of
-                    the passed `model`.
+                    projection method.
             train_dataset (torch.utils.data.Dataset or torch.utils.data.DataLoader):
                     In the `influence` method, we compute the influence score of
                     training examples on examples in a test batch.

captum/insights/attr_vis/server.py

@@ -4,7 +4,7 @@
 import socket
 import threading
 from time import sleep
-from typing import Optional
+from typing import cast, Dict, Optional
 
 from captum.log import log_usage
 from flask import Flask, jsonify, render_template, request
@@ -41,7 +41,7 @@ def namedtuple_to_dict(obj):
 def attribute() -> Response:
     # force=True needed for Colab notebooks, which doesn't use the correct
     # Content-Type header when forwarding requests through the Colab proxy
-    r = request.get_json(force=True)
+    r = cast(Dict, request.get_json(force=True))
     return jsonify(
         namedtuple_to_dict(
             visualizer._calculate_attribution_from_cache(  # type: ignore

setup.cfg

@@ -1,7 +1,7 @@
 [flake8]
 # E203: black and flake8 disagree on whitespace before ':'
 # W503: black and flake8 disagree on how to place operators
-ignore = E203, W503
+ignore = E203, W503, E704
 max-line-length = 88
 exclude =
   build, dist, tutorials, website

tests/attr/helpers/conductance_reference.py

@@ -1,5 +1,5 @@
 #!/usr/bin/env python3
-from typing import Optional, Tuple
+from typing import cast, Optional, Tuple, Union
 
 import numpy as np
 import torch
@@ -11,6 +11,7 @@
 from captum.attr._utils.attribution import LayerAttribution
 from captum.attr._utils.common import _reshape_and_sum
 from torch import Tensor
+from torch.utils.hooks import RemovableHandle
 
 """
 Note: This implementation of conductance follows the procedure described in the original
@@ -55,7 +56,7 @@ def forward_hook(module, inp, out):
             # The hidden layer tensor is assumed to have dimension (num_hidden, ...)
             # where the product of the dimensions >= 1 correspond to the total
             # number of hidden neurons in the layer.
-            layer_size = tuple(saved_tensor.size())[1:]
+            layer_size = tuple(cast(Tensor, saved_tensor).size())[1:]
             layer_units = int(np.prod(layer_size))
 
             # Remove unnecessary forward hook.
@@ -101,12 +102,12 @@ def forward_hook_register_back(module, inp, out):
             input_grads = torch.autograd.grad(torch.unbind(output), expanded_input)
 
             # Remove backwards hook
-            back_hook.remove()
+            cast(RemovableHandle, back_hook).remove()
 
             # Remove duplicates in gradient with respect to hidden layer,
             # choose one for each layer_units indices.
             output_mid_grads = torch.index_select(
-                saved_grads,
+                cast(Tensor, saved_grads),
                 0,
                 torch.tensor(range(0, input_grads[0].shape[0], layer_units)),
             )
@@ -115,7 +116,7 @@ def forward_hook_register_back(module, inp, out):
     def attribute(
         self,
         inputs,
-        baselines: Optional[int] = None,
+        baselines: Union[None, int, Tensor] = None,
         target=None,
         n_steps: int = 500,
         method: str = "riemann_trapezoid",

tests/attr/layer/test_layer_lrp.py

@@ -1,4 +1,5 @@
 #!/usr/bin/env python3
+# mypy: ignore-errors
 
 from typing import Any, Tuple
 

tests/attr/models/test_pytext.py

@@ -43,7 +43,7 @@ def __init__(self) -> None:
 
 
 class TestWordEmbeddings(unittest.TestCase):
-    def setUp(self) -> Optional[NoReturn]:
+    def setUp(self) -> None:
         if not HAS_PYTEXT:
             return self.skipTest("Skip the test since PyText is not installed")
 

tests/attr/test_class_summarizer.py

@@ -1,4 +1,6 @@
 #!/usr/bin/env python3
+from typing import List
+
 import torch
 from captum.attr import ClassSummarizer, CommonStats
 from tests.helpers.basic import BaseTest
@@ -45,7 +47,7 @@ def test_classes(self) -> None:
             ((3, 2, 10, 3), (1,)),
             # ((20,),),
         ]
-        list_of_classes = [
+        list_of_classes: List[List] = [
             list(range(100)),
             ["%d" % i for i in range(100)],
             list(range(300, 400)),

tests/attr/test_guided_grad_cam.py

@@ -1,7 +1,7 @@
 #!/usr/bin/env python3
 
 import unittest
-from typing import Any
+from typing import Any, List, Tuple, Union
 
 import torch
 from captum._utils.typing import TensorOrTupleOfTensorsGeneric
@@ -107,7 +107,7 @@ def _guided_grad_cam_test_assert(
         model: Module,
         target_layer: Module,
         test_input: TensorOrTupleOfTensorsGeneric,
-        expected: Tensor,
+        expected: Union[Tensor, List, Tuple],
         additional_input: Any = None,
         interpolate_mode: str = "nearest",
         attribute_to_layer_input: bool = False,

tests/attr/test_input_layer_wrapper.py

@@ -27,6 +27,7 @@
     BasicModel_MultiLayer_TrueMultiInput,
     MixedKwargsAndArgsModule,
 )
+from torch.nn import Module
 
 layer_methods_to_test_with_equiv = [
     # layer_method, equiv_method, whether or not to use multiple layers
@@ -115,7 +116,7 @@ def layer_method_with_input_layer_patches(
         assertTensorTuplesAlmostEqual(self, a1, real_attributions)
 
     def forward_eval_layer_with_inputs_helper(
-        self, model: ModelInputWrapper, inputs_to_test
+        self, model: Module, inputs_to_test
     ) -> None:
         # hard coding for simplicity
         # 0 if using args, 1 if using kwargs

tests/attr/test_lime.py

@@ -494,7 +494,7 @@ def _lime_test_assert(
         model: Callable,
         test_input: TensorOrTupleOfTensorsGeneric,
         expected_attr,
-        expected_coefs_only: Optional[Tensor] = None,
+        expected_coefs_only: Union[None, List, Tensor] = None,
         feature_mask: Union[None, TensorOrTupleOfTensorsGeneric] = None,
         additional_input: Any = None,
         perturbations_per_eval: Tuple[int, ...] = (1,),

tests/attr/test_stat.py

@@ -1,5 +1,6 @@
 #!/usr/bin/env python3
 import random
+from typing import Callable, List
 
 import torch
 from captum.attr import Max, Mean, Min, MSE, StdDev, Sum, Summarizer, Var
@@ -140,7 +141,7 @@ def test_stats_random_data(self) -> None:
             "sum",
             "mse",
         ]
-        gt_fns = [
+        gt_fns: List[Callable] = [
             torch.mean,
             lambda x: torch.var(x, unbiased=False),
             lambda x: torch.var(x, unbiased=True),

tests/helpers/basic.py

@@ -2,7 +2,7 @@
 import copy
 import random
 import unittest
-from typing import Callable
+from typing import Callable, List, Tuple, Union
 
 import numpy as np
 import torch
@@ -20,9 +20,7 @@ def copy_args(*args, **kwargs):
     return copy_args
 
 
-def assertTensorAlmostEqual(
-    test, actual: Tensor, expected: Tensor, delta: float = 0.0001, mode: str = "sum"
-) -> None:
+def assertTensorAlmostEqual(test, actual, expected, delta=0.0001, mode="sum"):
     assert isinstance(actual, torch.Tensor), (
         "Actual parameter given for " "comparison must be a tensor."
     )
@@ -60,9 +58,7 @@ def assertTensorAlmostEqual(
         raise ValueError("Mode for assertion comparison must be one of `max` or `sum`.")
 
 
-def assertTensorTuplesAlmostEqual(
-    test, actual, expected, delta: float = 0.0001, mode: str = "sum"
-) -> None:
+def assertTensorTuplesAlmostEqual(test, actual, expected, delta=0.0001, mode="sum"):
     if isinstance(expected, tuple):
         assert len(actual) == len(
             expected

tests/helpers/basic_models.py

@@ -44,7 +44,7 @@ def __init__(self) -> None:
         super().__init__()
 
     def forward(self, input: int):
-        input = 1 - F.relu(1 - input)
+        input = 1 - F.relu(torch.tensor(1 - input))
         return input
 
 

tests/influence/_core/test_tracin_validation.py

@@ -77,7 +77,7 @@ def test_tracincp_fast_rand_proj_inputs(self) -> None:
             ):
                 TracInCPFast(
                     net,
-                    "invalid_layer",
+                    "invalid_layer",  # type: ignore
                     train_dataset,
                     tmpdir,
                     loss_fn=nn.MSELoss(),

tests/robust/test_FGSM.py

@@ -188,7 +188,7 @@ def _FGSM_assert(
         inputs: TensorOrTupleOfTensorsGeneric,
         target: Any,
         epsilon: float,
-        answer: Union[TensorLikeList, Tuple[TensorLikeList, ...]],
+        answer: Union[List, Tuple[List, ...]],
         targeted: bool = False,
         additional_inputs: Any = None,
         lower_bound: float = float("-inf"),

tests/robust/test_attack_comparator.py

@@ -202,7 +202,7 @@ def test_attack_comparator_with_additional_args(self) -> None:
         attack_comp.reset()
         self.assertEqual(len(attack_comp.summary()), 0)
 
-    def _compare_results(self, obtained: Tensor, expected) -> None:
+    def _compare_results(self, obtained, expected) -> None:
         if isinstance(expected, dict):
             self.assertIsInstance(obtained, dict)
             for key in expected:

tests/utils/models/linear_models/_test_linear_classifier.py

@@ -1,6 +1,6 @@
 import argparse
 import random
-from typing import Optional
+from typing import cast, Optional
 
 import captum._utils.models.linear_model.model as pytorch_model_module
 import numpy as np
@@ -98,9 +98,9 @@ def compare_to_sk_learn(
         o_pytorch["l1_reg"] = alpha * pytorch_h.norm(p=1, dim=-1)
         o_sklearn["l1_reg"] = alpha * sklearn_h.norm(p=1, dim=-1)
 
-    rel_diff = (sum(o_sklearn.values()) - sum(o_pytorch.values())) / abs(
-        sum(o_sklearn.values())
-    )
+    rel_diff = cast(
+        np.ndarray, (sum(o_sklearn.values()) - sum(o_pytorch.values()))
+    ) / abs(sum(o_sklearn.values()))
     return (
         {
             "objective_rel_diff": rel_diff.tolist(),

tests/utils/test_helpers.py

@@ -7,7 +7,7 @@
 class HelpersTest(BaseTest):
     def test_assert_tensor_almost_equal(self) -> None:
         with self.assertRaises(AssertionError) as cm:
-            assertTensorAlmostEqual(self, [[1.0]], [[1.0]])
+            assertTensorAlmostEqual(self, [[1.0]], [[1.0]])  # type: ignore
         self.assertEqual(
             cm.exception.args,
             ("Actual parameter given for comparison must be a tensor.",),

tests/utils/test_linear_model.py

@@ -12,11 +12,11 @@
 from torch import Tensor
 
 
-def _evaluate(test_data, classifier) -> Dict[str, float]:
+def _evaluate(test_data, classifier) -> Dict[str, Tensor]:
     classifier.eval()
 
-    l1_loss = 0.0
-    l2_loss = 0.0
+    l1_loss = torch.tensor(0.0)
+    l2_loss = torch.tensor(0.0)
     n = 0
     l2_losses = []
     with torch.no_grad():
@@ -67,7 +67,7 @@ def train_and_compare(
         model_type,
         xs,
         ys,
-        expected_loss: Tensor,
+        expected_loss: Union[int, float, Tensor],
         expected_reg: Union[float, Tensor] = 0.0,
         expected_hyperplane: Optional[Tensor] = None,
         norm_hyperplane: bool = True,

tests/utils/test_sample_gradient.py

@@ -1,7 +1,7 @@
 #!/usr/bin/env python3
 
 import unittest
-from typing import Callable, Tuple
+from typing import Callable, List, Tuple
 
 import torch
 from captum._utils.gradient import apply_gradient_requirements
@@ -110,7 +110,11 @@ def test_sample_grads_layer_modules(self) -> None:
 
         # possible candidates for `layer_modules`, which are the modules whose
         # parameters we want to compute sample grads for
-        layer_moduless = [[model.conv1], [model.fc1], [model.conv1, model.fc1]]
+        layer_moduless: List[List[Module]] = [
+            [model.conv1],
+            [model.fc1],
+            [model.conv1, model.fc1],
+        ]
         # hard coded all modules we want to check
         all_modules = [model.conv1, model.fc1]
 
@@ -135,10 +139,10 @@ def test_sample_grads_layer_modules(self) -> None:
                     # So, check that we did calculate sample grads for the desired
                     # layers via the above checking approach.
                     for parameter in module.parameters():
-                        assert not isinstance(parameter.sample_grad, int)
+                        assert not isinstance(parameter.sample_grad, int)  # type: ignore
                 else:
                     # For the layers we do not want sample grads for, their
                     # `sample_grad` should still be 0, since they should not have been
                     # over-written.
                     for parameter in module.parameters():
-                        assert parameter.sample_grad == 0
+                        assert parameter.sample_grad == 0  # type: ignore