Fix Lime 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/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.
         """

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)

tests/attr/test_saliency.py

@@ -112,7 +112,7 @@ def test_futures_not_implemented(self) -> None:
         saliency = Saliency(model)
         attributions = None
         with self.assertRaises(NotImplementedError):
-            attributions = saliency.attribute_future()
+            attributions = saliency.attribute_future()  # type: ignore
         self.assertEqual(attributions, None)
 
     def _saliency_base_assert(