Testing linux job fix (meta-pytorch#1586)

Summary:
Pull Request resolved: meta-pytorch#1586

Rollback Plan:

Differential Revision: D76657713

Author: Vivek Miglani

.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/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/feature_ablation.py

@@ -401,10 +401,10 @@ def attribute(
         if attr_progress is not None:
             attr_progress.close()
 
-        # pyre-fixme[7]: Expected `Variable[TensorOrTupleOfTensorsGeneric <:
-        # [Tensor, typing.Tuple[Tensor, ...]]]`
-        # but got `Union[Tensor, typing.Tuple[Tensor, ...]]`.
-        return self._generate_result(total_attrib, weights, is_inputs_tuple)  # type: ignore # noqa: E501 line too long
+        return cast(
+            TensorOrTupleOfTensorsGeneric,
+            self._generate_result(total_attrib, weights, is_inputs_tuple),
+        )
 
     def _attribute_with_independent_feature_masks(
         self,
@@ -629,8 +629,7 @@ def _should_skip_inputs_and_warn(
                 all_empty = False
             if self._min_examples_per_batch_grouped is not None and (
                 formatted_inputs[tensor_idx].shape[0]
-                # pyre-ignore[58]: Type has been narrowed to int
-                < self._min_examples_per_batch_grouped
+                < cast(int, self._min_examples_per_batch_grouped)
             ):
                 should_skip = True
                 break
@@ -789,35 +788,35 @@ def attribute_future(
             )
 
             if enable_cross_tensor_attribution:
-                # pyre-fixme[7]: Expected `Future[Variable[TensorOrTupleOfTensorsGeneric
-                # <:[Tensor, typing.Tuple[Tensor, ...]]]]` but got
-                # `Future[Union[Tensor, typing.Tuple[Tensor, ...]]]`
-                return self._attribute_with_cross_tensor_feature_masks_future(  # type: ignore # noqa: E501 line too long
-                    formatted_inputs=formatted_inputs,
-                    formatted_additional_forward_args=formatted_additional_forward_args,
-                    target=target,
-                    baselines=baselines,
-                    formatted_feature_mask=formatted_feature_mask,
-                    attr_progress=attr_progress,
-                    processed_initial_eval_fut=processed_initial_eval_fut,
-                    is_inputs_tuple=is_inputs_tuple,
-                    perturbations_per_eval=perturbations_per_eval,
+                return cast(
+                    Future[TensorOrTupleOfTensorsGeneric],
+                    self._attribute_with_cross_tensor_feature_masks_future(  # type: ignore # noqa: E501 line too long
+                        formatted_inputs=formatted_inputs,
+                        formatted_additional_forward_args=formatted_additional_forward_args,
+                        target=target,
+                        baselines=baselines,
+                        formatted_feature_mask=formatted_feature_mask,
+                        attr_progress=attr_progress,
+                        processed_initial_eval_fut=processed_initial_eval_fut,
+                        is_inputs_tuple=is_inputs_tuple,
+                        perturbations_per_eval=perturbations_per_eval,
+                    ),
                 )
             else:
-                # pyre-fixme[7]: Expected `Future[Variable[TensorOrTupleOfTensorsGeneric
-                # <:[Tensor, typing.Tuple[Tensor, ...]]]]` but got
-                # `Future[Union[Tensor, typing.Tuple[Tensor, ...]]]`
-                return self._attribute_with_independent_feature_masks_future(  # type: ignore # noqa: E501 line too long
-                    formatted_inputs,
-                    formatted_additional_forward_args,
-                    target,
-                    baselines,
-                    formatted_feature_mask,
-                    perturbations_per_eval,
-                    attr_progress,
-                    processed_initial_eval_fut,
-                    is_inputs_tuple,
-                    **kwargs,
+                return cast(
+                    Future[TensorOrTupleOfTensorsGeneric],
+                    self._attribute_with_independent_feature_masks_future(  # type: ignore # noqa: E501 line too long
+                        formatted_inputs,
+                        formatted_additional_forward_args,
+                        target,
+                        baselines,
+                        formatted_feature_mask,
+                        perturbations_per_eval,
+                        attr_progress,
+                        processed_initial_eval_fut,
+                        is_inputs_tuple,
+                        **kwargs,
+                    ),
                 )
 
     def _attribute_with_independent_feature_masks_future(

captum/attr/_core/guided_backprop_deconvnet.py

@@ -2,7 +2,7 @@
 
 # pyre-strict
 import warnings
-from typing import Callable, List, Optional, Tuple, Union
+from typing import cast, List, Optional, Tuple, Union
 
 import torch
 import torch.nn.functional as F
@@ -78,12 +78,11 @@ def attribute(
             self._remove_hooks()
 
         undo_gradient_requirements(inputs_tuple, gradient_mask)
-        # pyre-fixme[7]: Expected `TensorOrTupleOfTensorsGeneric` but got
-        #  `Tuple[Tensor, ...]`.
-        return _format_output(is_inputs_tuple, gradients)
+        return cast(
+            TensorOrTupleOfTensorsGeneric, _format_output(is_inputs_tuple, gradients)
+        )
 
-    # 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 ModifiedReluGradientAttribution.
         """

captum/attr/_core/guided_grad_cam.py

@@ -2,7 +2,7 @@
 
 # pyre-strict
 import warnings
-from typing import List, Optional, Union
+from typing import cast, List, Optional, Union
 
 import torch
 from captum._utils.common import _format_output, _format_tensor_into_tuples, _is_tuple
@@ -223,6 +223,7 @@ def attribute(
                 )
                 output_attr.append(torch.empty(0))
 
-        # pyre-fixme[7]: Expected `TensorOrTupleOfTensorsGeneric` but got
-        #  `Tuple[Tensor, ...]`.
-        return _format_output(is_inputs_tuple, tuple(output_attr))
+        return cast(
+            TensorOrTupleOfTensorsGeneric,
+            _format_output(is_inputs_tuple, tuple(output_attr)),
+        )

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