Fix Black version and OSS Failures (#1241)

Summary:
Currently, OSS GitHub Actions tests are failing due to failing test, lint and typing issues. This updates the black version used externally (and corresponding python version to support latest black) to match the internal updates in D54447730 and also updates flake8 settings to avoid incompatibilities. Typing issues are also resolved and imports from torch._tensor are removed, since these are not supported for previous torch versions.

Pull Request resolved: #1241

Reviewed By: cyrjano

Differential Revision: D54901754

Pulled By: vivekmig

fbshipit-source-id: 2b94bf36488b11b6c145175cfe10fc5433b014fe

Author: vivekmig

.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."