Fix utils/model pyre fixme issues

.github/workflows/lint.yml

@@ -10,7 +10,7 @@ on:
 
 jobs:
   tests:
-    uses: pytorch/test-infra/.github/workflows/linux_job.yml@main
+    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main
     with:
       runner: linux.12xlarge
       docker-image: cimg/python:3.11

.github/workflows/test-pip-cpu-with-type-checks.yml

@@ -14,7 +14,7 @@ jobs:
       matrix:
         pytorch_args: ["", "-n"]
       fail-fast: false
-    uses: pytorch/test-infra/.github/workflows/linux_job.yml@main
+    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main
     with:
       runner: linux.12xlarge
       docker-image: cimg/python:3.11

.github/workflows/test-pip-cpu.yml

@@ -37,7 +37,7 @@ jobs:
           - pytorch_args: "-v 2.1.0"
             docker_img: "cimg/python:3.12"
       fail-fast: false
-    uses: pytorch/test-infra/.github/workflows/linux_job.yml@main
+    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main
     with:
       runner: linux.12xlarge
       docker-image: ${{ matrix.docker_img }}

.github/workflows/test-pip-gpu.yml

@@ -14,7 +14,7 @@ jobs:
       matrix:
         cuda_arch_version: ["12.1"]
       fail-fast: false
-    uses: pytorch/test-infra/.github/workflows/linux_job.yml@main
+    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main
     with:
       runner: linux.4xlarge.nvidia.gpu
       repository: pytorch/captum

captum/_utils/common.py

@@ -91,7 +91,7 @@ def _is_tuple(inputs: Tensor) -> Literal[False]: ...
 
 @typing.overload
 def _is_tuple(
-    inputs: TensorOrTupleOfTensorsGeneric,  # type: ignore
+    inputs: Union[Tensor, Tuple[Tensor, ...]],
 ) -> bool: ...
 
 
@@ -373,8 +373,6 @@ def _expand_target(
 def _expand_feature_mask(
     feature_mask: Union[Tensor, Tuple[Tensor, ...]], n_samples: int
 ) -> Tuple[Tensor, ...]:
-    # pyre-fixme[6]: For 1st argument expected `Tensor` but got `Union[Tensor,
-    #  typing.Tuple[Tensor, ...]]`.
     is_feature_mask_tuple = _is_tuple(feature_mask)
     feature_mask = _format_tensor_into_tuples(feature_mask)
     feature_mask_new = tuple(

captum/_utils/models/model.py

@@ -22,8 +22,7 @@ class Model(ABC):
     def fit(
         self,
         train_data: DataLoader,
-        # pyre-fixme[2]: Parameter must be annotated.
-        **kwargs,
+        **kwargs: object,
     ) -> Optional[Dict[str, Union[int, float, Tensor]]]:
         r"""
         Override this method to actually train your model.

captum/attr/_core/deep_lift.py

@@ -348,20 +348,21 @@ def attribute(  # type: ignore
             self._remove_hooks(main_model_hooks)
 
         undo_gradient_requirements(inputs_tuple, gradient_mask)
-        # pyre-fixme[7]: Expected `Union[Tuple[Variable[TensorOrTupleOfTensorsGeneric...
-        return _compute_conv_delta_and_format_attrs(
-            self,
-            return_convergence_delta,
-            attributions,
-            baselines,
-            inputs_tuple,
-            additional_forward_args,
-            target,
-            is_inputs_tuple,
+        return cast(
+            TensorOrTupleOfTensorsGeneric,
+            _compute_conv_delta_and_format_attrs(
+                self,
+                return_convergence_delta,
+                attributions,
+                baselines,
+                inputs_tuple,
+                additional_forward_args,
+                target,
+                is_inputs_tuple,
+            ),
         )
 
-    # pyre-fixme[24] Generic type `Callable` expects 2 type parameters.
-    def attribute_future(self) -> Callable:
+    def attribute_future(self) -> None:
         r"""
         This method is not implemented for DeepLift.
         """
@@ -831,11 +832,18 @@ def attribute(  # type: ignore
         )
 
         if return_convergence_delta:
-            # pyre-fixme[7]: Expected `Union[Tuple[Variable[TensorOrTupleOfTensorsGen...
-            return _format_output(is_inputs_tuple, attributions), delta
+            return (
+                cast(
+                    TensorOrTupleOfTensorsGeneric,
+                    _format_output(is_inputs_tuple, attributions),
+                ),
+                delta,
+            )
         else:
-            # pyre-fixme[7]: Expected `Union[Tuple[Variable[TensorOrTupleOfTensorsGen...
-            return _format_output(is_inputs_tuple, attributions)
+            return cast(
+                TensorOrTupleOfTensorsGeneric,
+                _format_output(is_inputs_tuple, attributions),
+            )
 
     def _expand_inputs_baselines_targets(
         self,
@@ -995,10 +1003,8 @@ def maxpool3d(
 
 def maxpool(
     module: Module,
-    # pyre-fixme[24]: Generic type `Callable` expects 2 type parameters.
-    pool_func: Callable,
-    # pyre-fixme[24]: Generic type `Callable` expects 2 type parameters.
-    unpool_func: Callable,
+    pool_func: Callable[..., Tuple[Tensor, Tensor]],
+    unpool_func: Callable[..., Tensor],
     inputs: Tensor,
     outputs: Tensor,
     grad_input: Tensor,

captum/attr/_core/input_x_gradient.py

@@ -1,7 +1,7 @@
 #!/usr/bin/env python3
 
 # pyre-strict
-from typing import Callable, Optional
+from typing import Callable, cast, Optional
 
 from captum._utils.common import _format_output, _format_tensor_into_tuples, _is_tuple
 from captum._utils.gradient import (
@@ -126,12 +126,11 @@ def attribute(
         )
 
         undo_gradient_requirements(inputs_tuple, gradient_mask)
-        # pyre-fixme[7]: Expected `TensorOrTupleOfTensorsGeneric` but got
-        #  `Tuple[Tensor, ...]`.
-        return _format_output(is_inputs_tuple, attributions)
+        return cast(
+            TensorOrTupleOfTensorsGeneric, _format_output(is_inputs_tuple, attributions)
+        )
 
-    # pyre-fixme[24] Generic type `Callable` expects 2 type parameters.
-    def attribute_future(self) -> Callable:
+    def attribute_future(self) -> None:
         r"""
         This method is not implemented for InputXGradient.
         """

captum/attr/_core/integrated_gradients.py

@@ -2,7 +2,7 @@
 
 # pyre-strict
 import typing
-from typing import Callable, List, Literal, Optional, Tuple, Union
+from typing import Callable, cast, List, Literal, Optional, Tuple, Union
 
 import torch
 from captum._utils.common import (
@@ -301,16 +301,18 @@ def attribute(  # type: ignore
                 additional_forward_args=additional_forward_args,
                 target=target,
             )
-            # pyre-fixme[7]: Expected `Union[Tuple[Variable[TensorOrTupleOfTensorsGen...
-            return _format_output(is_inputs_tuple, attributions), delta
-        # pyre-fixme[7]: Expected
-        #  `Union[Tuple[Variable[TensorOrTupleOfTensorsGeneric <: [Tensor,
-        #  typing.Tuple[Tensor, ...]]], Tensor], Variable[TensorOrTupleOfTensorsGeneric
-        #  <: [Tensor, typing.Tuple[Tensor, ...]]]]` but got `Tuple[Tensor, ...]`.
-        return _format_output(is_inputs_tuple, attributions)
-
-    # pyre-fixme[24] Generic type `Callable` expects 2 type parameters.
-    def attribute_future(self) -> Callable:
+            return (
+                cast(
+                    TensorOrTupleOfTensorsGeneric,
+                    _format_output(is_inputs_tuple, attributions),
+                ),
+                delta,
+            )
+        return cast(
+            TensorOrTupleOfTensorsGeneric, _format_output(is_inputs_tuple, attributions)
+        )
+
+    def attribute_future(self) -> None:
         r"""
         This method is not implemented for IntegratedGradients.
         """

captum/attr/_core/kernel_shap.py

@@ -292,8 +292,7 @@ def attribute(  # type: ignore
             show_progress=show_progress,
         )
 
-    # pyre-fixme[24] Generic type `Callable` expects 2 type parameters.
-    def attribute_future(self) -> Callable:
+    def attribute_future(self) -> None:
         r"""
         This method is not implemented for KernelShap.
         """

captum/attr/_core/lime.py

@@ -548,8 +548,7 @@ def generate_perturbation() -> (
 
         return generate_perturbation
 
-    # pyre-fixme[24] Generic type `Callable` expects 2 type parameters.
-    def attribute_future(self) -> Callable:
+    def attribute_future(self) -> None:
         r"""
         This method is not implemented for LimeBase.
         """
@@ -1116,8 +1115,7 @@ def attribute(  # type: ignore
             show_progress=show_progress,
         )
 
-    # pyre-fixme[24] Generic type `Callable` expects 2 type parameters.
-    def attribute_future(self) -> Callable:
+    def attribute_future(self) -> None:
         return super().attribute_future()
 
     def _attribute_kwargs(  # type: ignore

captum/attr/_core/lrp.py

@@ -4,7 +4,7 @@
 
 import typing
 from collections import defaultdict
-from typing import Any, Callable, cast, Dict, List, Literal, Optional, Tuple, Union
+from typing import Any, cast, Dict, List, Literal, Optional, Tuple, Union
 
 import torch.nn as nn
 from captum._utils.common import (
@@ -230,16 +230,17 @@ def attribute(
         undo_gradient_requirements(input_tuple, gradient_mask)
 
         if return_convergence_delta:
-            # pyre-fixme[7]: Expected `Union[Tuple[Variable[TensorOrTupleOfTensorsGen...
             return (
-                _format_output(is_inputs_tuple, relevances),
+                cast(
+                    TensorOrTupleOfTensorsGeneric,
+                    _format_output(is_inputs_tuple, relevances),
+                ),
                 self.compute_convergence_delta(relevances, output),
             )
         else:
             return _format_output(is_inputs_tuple, relevances)  # type: ignore
 
-    # pyre-fixme[24] Generic type `Callable` expects 2 type parameters.
-    def attribute_future(self) -> Callable:
+    def attribute_future(self) -> None:
         r"""
         This method is not implemented for LRP.
         """

captum/attr/_core/occlusion.py

@@ -269,8 +269,7 @@ def attribute(  # type: ignore
             show_progress=show_progress,
         )
 
-    # pyre-fixme[24] Generic type `Callable` expects 2 type parameters.
-    def attribute_future(self) -> Callable:
+    def attribute_future(self) -> None:
         r"""
         This method is not implemented for Occlusion.
         """

captum/attr/_core/saliency.py

@@ -2,7 +2,7 @@
 
 # pyre-strict
 
-from typing import Callable, Optional
+from typing import Callable, cast, Optional
 
 import torch
 from captum._utils.common import _format_output, _format_tensor_into_tuples, _is_tuple
@@ -138,12 +138,12 @@ def attribute(
         else:
             attributions = gradients
         undo_gradient_requirements(inputs_tuple, gradient_mask)
-        # pyre-fixme[7]: Expected `TensorOrTupleOfTensorsGeneric` but got
-        #  `Tuple[Tensor, ...]`.
-        return _format_output(is_inputs_tuple, attributions)
 
-    # pyre-fixme[24] Generic type `Callable` expects 2 type parameters.
-    def attribute_future(self) -> Callable:
+        return cast(
+            TensorOrTupleOfTensorsGeneric, _format_output(is_inputs_tuple, attributions)
+        )
+
+    def attribute_future(self) -> None:
         r"""
         This method is not implemented for Saliency.
         """

captum/testing/attr/helpers/get_config_util.py

@@ -1,7 +1,7 @@
 # (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary.
 
 # pyre-strict
-from typing import Any, Tuple
+from typing import List, Tuple
 
 import torch
 from captum._utils.gradient import compute_gradients
@@ -10,17 +10,15 @@
 from torch.nn import Module
 
 
-# pyre-fixme[3]: Return annotation cannot contain `Any`.
-def get_basic_config() -> Tuple[Module, Tensor, Tensor, Any]:
+def get_basic_config() -> Tuple[Module, Tensor, Tensor, None]:
     input = torch.tensor([1.0, 2.0, 3.0, 0.0, -1.0, 7.0], requires_grad=True).T
     # manually percomputed gradients
     grads = torch.tensor([-0.0, -0.0, -0.0, 1.0, 1.0, -0.0])
     return BasicModel(), input, grads, None
 
 
-# pyre-fixme[3]: Return annotation cannot contain `Any`.
 def get_multiargs_basic_config() -> (
-    Tuple[Module, Tuple[Tensor, ...], Tuple[Tensor, ...], Any]
+    Tuple[Module, Tuple[Tensor, ...], Tuple[Tensor, ...], Tuple[List[int], int]]
 ):
     model = BasicModel5_MultiArgs()
     additional_forward_args = ([2, 3], 1)
@@ -34,9 +32,8 @@ def get_multiargs_basic_config() -> (
     return model, inputs, grads, additional_forward_args
 
 
-# pyre-fixme[3]: Return annotation cannot contain `Any`.
 def get_multiargs_basic_config_large() -> (
-    Tuple[Module, Tuple[Tensor, ...], Tuple[Tensor, ...], Any]
+    Tuple[Module, Tuple[Tensor, ...], Tuple[Tensor, ...], Tuple[List[int], int]]
 ):
     model = BasicModel5_MultiArgs()
     additional_forward_args = ([2, 3], 1)

tests/attr/test_deeplift_basic.py

@@ -310,7 +310,7 @@ def test_futures_not_implemented(self) -> None:
         dl = DeepLift(model, multiply_by_inputs=False)
         attributions = None
         with self.assertRaises(NotImplementedError):
-            attributions = dl.attribute_future()
+            attributions = dl.attribute_future()  # type: ignore
         self.assertEqual(attributions, None)
 
     def _deeplift_assert(

tests/attr/test_input_x_gradient.py

@@ -55,7 +55,7 @@ def test_futures_not_implemented(self) -> None:
         input_x_grad = InputXGradient(model.forward)
         attributions = None
         with self.assertRaises(NotImplementedError):
-            attributions = input_x_grad.attribute_future()
+            attributions = input_x_grad.attribute_future()  # type: ignore
         self.assertEqual(attributions, None)
 
     def _input_x_gradient_base_assert(

tests/attr/test_integrated_gradients_basic.py

@@ -161,7 +161,7 @@ def test_futures_not_implemented(self) -> None:
         ig = IntegratedGradients(model, multiply_by_inputs=True)
         attributions = None
         with self.assertRaises(NotImplementedError):
-            attributions = ig.attribute_future()
+            attributions = ig.attribute_future()  # type: ignore
         self.assertEqual(attributions, None)
 
     def _assert_multi_variable(

tests/attr/test_kernel_shap.py

@@ -337,7 +337,7 @@ def test_futures_not_implemented(self) -> None:
         kernel_shap = KernelShap(net)
         attributions = None
         with self.assertRaises(NotImplementedError):
-            attributions = kernel_shap.attribute_future()
+            attributions = kernel_shap.attribute_future()  # type: ignore
         self.assertEqual(attributions, None)
 
     def _multi_input_scalar_kernel_shap_assert(self, func: Callable) -> None:

tests/attr/test_lime.py

@@ -482,7 +482,7 @@ def test_futures_not_implemented(self) -> None:
         )
         attributions = None
         with self.assertRaises(NotImplementedError):
-            attributions = lime.attribute_future()
+            attributions = lime.attribute_future()  # type: ignore
         self.assertEqual(attributions, None)
 
     # pyre-fixme[24]: Generic type `Callable` expects 2 type parameters.

tests/attr/test_lrp.py

@@ -335,5 +335,5 @@ def test_futures_not_implemented(self) -> None:
         lrp = LRP(model)
         attributions = None
         with self.assertRaises(NotImplementedError):
-            attributions = lrp.attribute_future()
+            attributions = lrp.attribute_future()  # type: ignore
         self.assertEqual(attributions, None)

tests/attr/test_occlusion.py

@@ -287,7 +287,7 @@ def test_futures_not_implemented(self) -> None:
         occ = Occlusion(net)
         attributions = None
         with self.assertRaises(NotImplementedError):
-            attributions = occ.attribute_future()
+            attributions = occ.attribute_future()  # type: ignore
         self.assertEqual(attributions, None)
 
     @unittest.mock.patch("sys.stderr", new_callable=io.StringIO)