Fix Black version and OSS Failures

.github/workflows/lint.yml

@@ -13,7 +13,7 @@ jobs:
     uses: pytorch/test-infra/.github/workflows/linux_job.yml@main
     with:
       runner: linux.12xlarge
-      docker-image: cimg/python:3.6
+      docker-image: cimg/python:3.9
       repository: pytorch/captum
       script: |
         sudo chmod -R 777 .

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/influence/_core/tracincp_fast_rand_proj.py

@@ -82,7 +82,7 @@ class TracInCPFast(TracInCPBase):
     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,
@@ -96,11 +96,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.
@@ -869,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,
@@ -886,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/attribution_calculation.py

@@ -131,7 +131,7 @@ def calculate_attribution(
         )
         if "baselines" in inspect.signature(attribution_method.attribute).parameters:
             attribution_arguments["baselines"] = baseline
-        attr = attribution_method.attribute.__wrapped__(
+        attr = attribution_method.attribute.__wrapped__(  # type: ignore
             attribution_method,  # self
             data,
             additional_forward_args=additional_forward_args,

captum/insights/attr_vis/features.py

@@ -8,7 +8,7 @@
 from captum._utils.common import safe_div
 from captum.attr._utils import visualization as viz
 from captum.insights.attr_vis._utils.transforms import format_transforms
-from torch._tensor import Tensor
+from torch import Tensor
 
 FeatureOutput = namedtuple("FeatureOutput", "name base modified type contribution")
 

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,10 +41,10 @@ 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(
+            visualizer._calculate_attribution_from_cache(  # type: ignore
                 r["inputIndex"], r["modelIndex"], r["labelIndex"]
             )
         )
@@ -54,15 +54,15 @@ def attribute() -> Response:
 @app.route("/fetch", methods=["POST"])
 def fetch() -> Response:
     # force=True needed, see comment for "/attribute" route above
-    visualizer._update_config(request.get_json(force=True))
-    visualizer_output = visualizer.visualize()
+    visualizer._update_config(request.get_json(force=True))  # type: ignore
+    visualizer_output = visualizer.visualize()  # type: ignore
     clean_output = namedtuple_to_dict(visualizer_output)
     return jsonify(clean_output)
 
 
 @app.route("/init")
 def init() -> Response:
-    return jsonify(visualizer.get_insights_config())
+    return jsonify(visualizer.get_insights_config())  # type: ignore
 
 
 @app.route("/")

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

setup.py

@@ -67,7 +67,7 @@ def report(*args):
     INSIGHTS_REQUIRES
     + TEST_REQUIRES
     + [
-        "black==22.3.0",
+        "black",
         "flake8",
         "sphinx",
         "sphinx-autodoc-typehints",

tests/attr/helpers/conductance_reference.py

@@ -1,5 +1,5 @@
 #!/usr/bin/env python3
-from typing import Optional, Tuple
+from typing import cast, Tuple, Union
 
 import numpy as np
 import torch
@@ -10,7 +10,8 @@
 from captum.attr._utils.approximation_methods import approximation_parameters
 from captum.attr._utils.attribution import LayerAttribution
 from captum.attr._utils.common import _reshape_and_sum
-from torch._tensor import Tensor
+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
 
@@ -9,7 +10,7 @@
 
 from tests.helpers.basic import assertTensorAlmostEqual, BaseTest
 from tests.helpers.basic_models import BasicModel_ConvNet_One_Conv, SimpleLRPModel
-from torch._tensor import Tensor
+from torch import Tensor
 
 
 def _get_basic_config() -> Tuple[BasicModel_ConvNet_One_Conv, Tensor]:

tests/attr/models/test_pytext.py

@@ -5,10 +5,9 @@
 import os
 import tempfile
 import unittest
-from typing import Dict, List, NoReturn, Optional
+from typing import Dict, List
 
 import torch
-from pytext.data.data_handler import CommonMetadata
 
 HAS_PYTEXT = True
 try:
@@ -20,6 +19,7 @@
     from pytext.config.component import create_featurizer, create_model
     from pytext.config.doc_classification import ModelInputConfig, TargetConfig
     from pytext.config.field_config import FeatureConfig, WordFeatConfig
+    from pytext.data.data_handler import CommonMetadata
     from pytext.data.doc_classification_data_handler import (  # @manual=//pytext:main_lib  # noqa
         DocClassificationDataHandler,
     )
@@ -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")
 
@@ -143,7 +143,7 @@ def _create_dummy_model(self):
             self._create_dummy_meta_data(),
         )
 
-    def _create_dummy_meta_data(self) -> CommonMetadata:
+    def _create_dummy_meta_data(self):
         text_field_meta = FieldMeta()
         text_field_meta.vocab = VocabStub()
         text_field_meta.vocab_size = 4

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_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/attr/test_guided_grad_cam.py

@@ -1,14 +1,14 @@
 #!/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
 from captum.attr._core.guided_grad_cam import GuidedGradCam
 from tests.helpers.basic import assertTensorAlmostEqual, BaseTest
 from tests.helpers.basic_models import BasicModel_ConvNet_One_Conv
-from torch._tensor import Tensor
+from torch import Tensor
 from torch.nn import Module
 
 
@@ -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_interpretable_input.py

@@ -4,7 +4,7 @@
 from captum.attr._utils.interpretable_input import TextTemplateInput, TextTokenInput
 from parameterized import parameterized
 from tests.helpers.basic import assertTensorAlmostEqual, BaseTest
-from torch._tensor import Tensor
+from torch import Tensor
 
 
 class DummyTokenizer:

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/concept/test_tcav.py

@@ -14,7 +14,6 @@
     Iterator,
     List,
     Set,
-    SupportsIndex,
     Tuple,
     Union,
 )
@@ -174,7 +173,7 @@ def __init__(
         self,
         get_tensor_from_filename_func: Callable,
         path: str,
-        num_samples: SupportsIndex = 100,
+        num_samples: int = 100,
     ) -> None:
         r"""
         Args:

tests/helpers/basic.py

@@ -7,7 +7,6 @@
 import numpy as np
 import torch
 from captum.log import patch_methods
-from torch._tensor import Tensor
 
 
 def deep_copy_args(func: Callable):
@@ -21,7 +20,7 @@ def copy_args(*args, **kwargs):
 
 
 def assertTensorAlmostEqual(
-    test, actual: Tensor, expected: Tensor, delta: float = 0.0001, mode: str = "sum"
+    test, actual, expected, delta: float = 0.0001, mode: str = "sum"
 ) -> None:
     assert isinstance(actual, torch.Tensor), (
         "Actual parameter given for " "comparison must be a tensor."