Improving typing of additional_forward_args

Summary: Change `additional_forward_args` to resolve pyre errors and address the feedback on D64998803. Differential Revision: D65178564
pytorch · Oct 30, 2024 · 1a4e012 · 1a4e012
1 parent 07470af
commit 1a4e012
Show file tree

Hide file tree

Showing 48 changed files with 172 additions and 210 deletions.
diff --git a/captum/_utils/av.py b/captum/_utils/av.py
@@ -351,8 +351,7 @@ def _compute_and_save_activations(
         inputs: Union[Tensor, Tuple[Tensor, ...]],
         identifier: str,
         num_id: str,
-        # pyre-fixme[2]: Parameter annotation cannot be `Any`.
-        additional_forward_args: Any = None,
+        additional_forward_args: Optional[object] = None,
         load_from_disk: bool = True,
     ) -> None:
         r"""

diff --git a/captum/_utils/common.py b/captum/_utils/common.py
@@ -273,9 +273,25 @@ def _format_float_or_tensor_into_tuples(
     return inputs
 
 
+@overload
+def _format_additional_forward_args(
+    # pyre-fixme[24]: Generic type `tuple` expects at least 1 type parameter.
+    additional_forward_args: Union[Tensor, Tuple]
+    # pyre-fixme[24]: Generic type `tuple` expects at least 1 type parameter.
+) -> Tuple: ...
+
+
+@overload
+def _format_additional_forward_args(  # type: ignore
+    additional_forward_args: Optional[object],
+    # pyre-fixme[24]: Generic type `tuple` expects at least 1 type parameter.
+) -> Union[None, Tuple]: ...
+
+
 def _format_additional_forward_args(
     additional_forward_args: Optional[object],
-) -> Union[None, Tuple[object, ...]]:
+    # pyre-fixme[24]: Generic type `tuple` expects at least 1 type parameter.
+) -> Union[None, Tuple]:
     if additional_forward_args is not None and not isinstance(
         additional_forward_args, tuple
     ):
@@ -284,8 +300,8 @@ def _format_additional_forward_args(
 
 
 def _expand_additional_forward_args(
-    # pyre-fixme[2]: Parameter annotation cannot be `Any`.
-    additional_forward_args: Any,
+    # pyre-fixme[24]: Generic type `tuple` expects at least 1 type parameter.
+    additional_forward_args: Union[None, Tuple],
     n_steps: int,
     expansion_type: ExpansionTypes = ExpansionTypes.repeat,
     # pyre-fixme[24]: Generic type `tuple` expects at least 1 type parameter.
@@ -557,8 +573,7 @@ def _run_forward(
     # pyre-fixme[2]: Parameter annotation cannot be `Any`.
     inputs: Any,
     target: TargetType = None,
-    # pyre-fixme[2]: Parameter annotation cannot be `Any`.
-    additional_forward_args: Any = None,
+    additional_forward_args: Optional[object] = None,
 ) -> Union[Tensor, Future[Tensor]]:
     forward_func_args = signature(forward_func).parameters
     if len(forward_func_args) == 0:

diff --git a/captum/_utils/gradient.py b/captum/_utils/gradient.py
@@ -104,8 +104,7 @@ def compute_gradients(
     forward_fn: Callable,
     inputs: Union[Tensor, Tuple[Tensor, ...]],
     target_ind: TargetType = None,
-    # pyre-fixme[2]: Parameter annotation cannot be `Any`.
-    additional_forward_args: Any = None,
+    additional_forward_args: Optional[object] = None,
 ) -> Tuple[Tensor, ...]:
     r"""
     Computes gradients of the output with respect to inputs for an
@@ -175,8 +174,7 @@ def _forward_layer_eval(
     forward_fn: Callable,
     inputs: Union[Tensor, Tuple[Tensor, ...]],
     layer: List[Module],
-    # pyre-fixme[2]: Parameter annotation cannot be `Any`.
-    additional_forward_args: Any = None,
+    additional_forward_args: Optional[object] = None,
     device_ids: Union[None, List[int]] = None,
     attribute_to_layer_input: bool = False,
     grad_enabled: bool = False,
@@ -191,8 +189,7 @@ def _forward_layer_eval(
     forward_fn: Callable,
     inputs: Union[Tensor, Tuple[Tensor, ...]],
     layer: Module,
-    # pyre-fixme[2]: Parameter annotation cannot be `Any`.
-    additional_forward_args: Any = None,
+    additional_forward_args: Optional[object] = None,
     device_ids: Union[None, List[int]] = None,
     attribute_to_layer_input: bool = False,
     grad_enabled: bool = False,
@@ -204,7 +201,7 @@ def _forward_layer_eval(
     forward_fn: Callable,
     inputs: Union[Tensor, Tuple[Tensor, ...]],
     layer: ModuleOrModuleList,
-    additional_forward_args: Any = None,
+    additional_forward_args: Optional[object] = None,
     device_ids: Union[None, List[int]] = None,
     attribute_to_layer_input: bool = False,
     grad_enabled: bool = False,
@@ -233,8 +230,7 @@ def _forward_layer_distributed_eval(
     inputs: Any,
     layer: ModuleOrModuleList,
     target_ind: TargetType = None,
-    # pyre-fixme[2]: Parameter annotation cannot be `Any`.
-    additional_forward_args: Any = None,
+    additional_forward_args: Optional[object] = None,
     attribute_to_layer_input: bool = False,
     forward_hook_with_return: Literal[False] = False,
     require_layer_grads: bool = False,
@@ -250,7 +246,7 @@ def _forward_layer_distributed_eval(
     inputs: Any,
     layer: ModuleOrModuleList,
     target_ind: TargetType = None,
-    additional_forward_args: Any = None,
+    additional_forward_args: Optional[object] = None,
     attribute_to_layer_input: bool = False,
     *,
     forward_hook_with_return: Literal[True],
@@ -264,7 +260,7 @@ def _forward_layer_distributed_eval(
     inputs: Any,
     layer: ModuleOrModuleList,
     target_ind: TargetType = None,
-    additional_forward_args: Any = None,
+    additional_forward_args: Optional[object] = None,
     attribute_to_layer_input: bool = False,
     forward_hook_with_return: bool = False,
     require_layer_grads: bool = False,
@@ -427,8 +423,7 @@ def _forward_layer_eval_with_neuron_grads(
     forward_fn: Callable,
     inputs: Union[Tensor, Tuple[Tensor, ...]],
     layer: Module,
-    # pyre-fixme[2]: Parameter annotation cannot be `Any`.
-    additional_forward_args: Any = None,
+    additional_forward_args: Optional[object] = None,
     *,
     # pyre-fixme[24]: Generic type `Callable` expects 2 type parameters.
     gradient_neuron_selector: Union[int, Tuple[Union[int, slice], ...], Callable],
@@ -446,7 +441,7 @@ def _forward_layer_eval_with_neuron_grads(
     forward_fn: Callable,
     inputs: Union[Tensor, Tuple[Tensor, ...]],
     layer: List[Module],
-    additional_forward_args: Any = None,
+    additional_forward_args: Optional[object] = None,
     gradient_neuron_selector: None = None,
     grad_enabled: bool = False,
     device_ids: Union[None, List[int]] = None,
@@ -462,7 +457,7 @@ def _forward_layer_eval_with_neuron_grads(
     forward_fn: Callable,
     inputs: Union[Tensor, Tuple[Tensor, ...]],
     layer: Module,
-    additional_forward_args: Any = None,
+    additional_forward_args: Optional[object] = None,
     gradient_neuron_selector: None = None,
     grad_enabled: bool = False,
     device_ids: Union[None, List[int]] = None,
@@ -475,7 +470,7 @@ def _forward_layer_eval_with_neuron_grads(
     forward_fn: Callable,
     inputs: Union[Tensor, Tuple[Tensor, ...]],
     layer: ModuleOrModuleList,
-    additional_forward_args: Any = None,
+    additional_forward_args: Optional[object] = None,
     # pyre-fixme[24]: Generic type `Callable` expects 2 type parameters.
     gradient_neuron_selector: Union[
         None, int, Tuple[Union[int, slice], ...], Callable
@@ -549,8 +544,7 @@ def compute_layer_gradients_and_eval(
     layer: Module,
     inputs: Union[Tensor, Tuple[Tensor, ...]],
     target_ind: TargetType = None,
-    # pyre-fixme[2]: Parameter annotation cannot be `Any`.
-    additional_forward_args: Any = None,
+    additional_forward_args: Optional[object] = None,
     *,
     # pyre-fixme[24]: Generic type `Callable` expects 2 type parameters.
     gradient_neuron_selector: Union[int, Tuple[Union[int, slice], ...], Callable],
@@ -571,7 +565,7 @@ def compute_layer_gradients_and_eval(
     layer: List[Module],
     inputs: Union[Tensor, Tuple[Tensor, ...]],
     target_ind: TargetType = None,
-    additional_forward_args: Any = None,
+    additional_forward_args: Optional[object] = None,
     gradient_neuron_selector: None = None,
     device_ids: Union[None, List[int]] = None,
     attribute_to_layer_input: bool = False,
@@ -590,7 +584,7 @@ def compute_layer_gradients_and_eval(
     layer: Module,
     inputs: Union[Tensor, Tuple[Tensor, ...]],
     target_ind: TargetType = None,
-    additional_forward_args: Any = None,
+    additional_forward_args: Optional[object] = None,
     gradient_neuron_selector: None = None,
     device_ids: Union[None, List[int]] = None,
     attribute_to_layer_input: bool = False,
@@ -606,7 +600,7 @@ def compute_layer_gradients_and_eval(
     layer: ModuleOrModuleList,
     inputs: Union[Tensor, Tuple[Tensor, ...]],
     target_ind: TargetType = None,
-    additional_forward_args: Any = None,
+    additional_forward_args: Optional[object] = None,
     # pyre-fixme[24]: Generic type `Callable` expects 2 type parameters.
     gradient_neuron_selector: Union[
         None, int, Tuple[Union[int, slice], ...], Callable
@@ -792,8 +786,7 @@ def grad_fn(
         forward_fn: Callable,
         inputs: TensorOrTupleOfTensorsGeneric,
         target_ind: TargetType = None,
-        # pyre-fixme[2]: Parameter annotation cannot be `Any`.
-        additional_forward_args: Any = None,
+        additional_forward_args: Optional[object] = None,
     ) -> Tuple[Tensor, ...]:
         _, grads = _forward_layer_eval_with_neuron_grads(
             forward_fn,

diff --git a/captum/attr/_core/deep_lift.py b/captum/attr/_core/deep_lift.py
@@ -3,7 +3,7 @@
 # pyre-strict
 import typing
 import warnings
-from typing import Callable, cast, Dict, List, Literal, Tuple, Type, Union
+from typing import Callable, cast, Dict, List, Literal, Optional, Tuple, Type, Union
 
 import torch
 import torch.nn as nn
@@ -117,7 +117,7 @@ def attribute(
         inputs: TensorOrTupleOfTensorsGeneric,
         baselines: BaselineType = None,
         target: TargetType = None,
-        additional_forward_args: object = None,
+        additional_forward_args: Optional[Tuple[object, ...]] = None,
         *,
         return_convergence_delta: Literal[True],
         custom_attribution_func: Union[None, Callable[..., Tuple[Tensor, ...]]] = None,
@@ -129,7 +129,7 @@ def attribute(
         inputs: TensorOrTupleOfTensorsGeneric,
         baselines: BaselineType = None,
         target: TargetType = None,
-        additional_forward_args: object = None,
+        additional_forward_args: Optional[Tuple[object, ...]] = None,
         return_convergence_delta: Literal[False] = False,
         custom_attribution_func: Union[None, Callable[..., Tuple[Tensor, ...]]] = None,
     ) -> TensorOrTupleOfTensorsGeneric: ...
@@ -140,7 +140,7 @@ def attribute(  # type: ignore
         inputs: TensorOrTupleOfTensorsGeneric,
         baselines: BaselineType = None,
         target: TargetType = None,
-        additional_forward_args: object = None,
+        additional_forward_args: Optional[Tuple[object, ...]] = None,
         return_convergence_delta: bool = False,
         custom_attribution_func: Union[None, Callable[..., Tuple[Tensor, ...]]] = None,
     ) -> Union[
@@ -370,7 +370,7 @@ def _construct_forward_func(
         forward_func: Callable[..., Tensor],
         inputs: Tuple[Tuple[Tensor, ...], Tuple[Tensor, ...]],
         target: TargetType = None,
-        additional_forward_args: object = None,
+        additional_forward_args: Optional[Tuple[object, ...]] = None,
     ) -> Callable[[], Tensor]:
         def forward_fn() -> Tensor:
             model_out = cast(
@@ -604,7 +604,7 @@ def attribute(
             TensorOrTupleOfTensorsGeneric, Callable[..., TensorOrTupleOfTensorsGeneric]
         ],
         target: TargetType = None,
-        additional_forward_args: object = None,
+        additional_forward_args: Optional[Tuple[object, ...]] = None,
         *,
         return_convergence_delta: Literal[True],
         custom_attribution_func: Union[None, Callable[..., Tuple[Tensor, ...]]] = None,
@@ -618,7 +618,7 @@ def attribute(
             TensorOrTupleOfTensorsGeneric, Callable[..., TensorOrTupleOfTensorsGeneric]
         ],
         target: TargetType = None,
-        additional_forward_args: object = None,
+        additional_forward_args: Optional[Tuple[object, ...]] = None,
         return_convergence_delta: Literal[False] = False,
         custom_attribution_func: Union[None, Callable[..., Tuple[Tensor, ...]]] = None,
     ) -> TensorOrTupleOfTensorsGeneric: ...
@@ -631,7 +631,7 @@ def attribute(  # type: ignore
             TensorOrTupleOfTensorsGeneric, Callable[..., TensorOrTupleOfTensorsGeneric]
         ],
         target: TargetType = None,
-        additional_forward_args: object = None,
+        additional_forward_args: Optional[Tuple[object, ...]] = None,
         return_convergence_delta: bool = False,
         custom_attribution_func: Union[None, Callable[..., Tuple[Tensor, ...]]] = None,
     ) -> Union[
@@ -840,7 +840,7 @@ def _expand_inputs_baselines_targets(
         baselines: Tuple[Tensor, ...],
         inputs: Tuple[Tensor, ...],
         target: TargetType,
-        additional_forward_args: object,
+        additional_forward_args: Optional[Tuple[object, ...]],
     ) -> Tuple[Tuple[Tensor, ...], Tuple[Tensor, ...], TargetType, object]:
         inp_bsz = inputs[0].shape[0]
         base_bsz = baselines[0].shape[0]

diff --git a/captum/attr/_core/feature_ablation.py b/captum/attr/_core/feature_ablation.py
@@ -75,7 +75,7 @@ def attribute(
         inputs: TensorOrTupleOfTensorsGeneric,
         baselines: BaselineType = None,
         target: TargetType = None,
-        additional_forward_args: object = None,
+        additional_forward_args: Optional[object] = None,
         feature_mask: Union[None, Tensor, Tuple[Tensor, ...]] = None,
         perturbations_per_eval: int = 1,
         show_progress: bool = False,
@@ -408,7 +408,7 @@ def attribute_future(
         inputs: TensorOrTupleOfTensorsGeneric,
         baselines: BaselineType = None,
         target: TargetType = None,
-        additional_forward_args: object = None,
+        additional_forward_args: Optional[object] = None,
         feature_mask: Union[None, Tensor, Tuple[Tensor, ...]] = None,
         perturbations_per_eval: int = 1,
         show_progress: bool = False,
@@ -655,7 +655,7 @@ def _ith_input_ablation_generator(
         self,
         i: int,
         inputs: TensorOrTupleOfTensorsGeneric,
-        additional_args: object,
+        additional_args: Optional[Tuple[object, ...]],
         target: TargetType,
         baselines: BaselineType,
         input_mask: Union[None, Tensor, Tuple[Tensor, ...]],

diff --git a/captum/attr/_core/feature_permutation.py b/captum/attr/_core/feature_permutation.py
@@ -1,7 +1,7 @@
 #!/usr/bin/env python3
 
 # pyre-strict
-from typing import Any, Callable, Tuple, Union
+from typing import Any, Callable, Optional, Tuple, Union
 
 import torch
 from captum._utils.typing import TargetType, TensorOrTupleOfTensorsGeneric
@@ -99,7 +99,7 @@ def attribute(  # type: ignore
         self,
         inputs: TensorOrTupleOfTensorsGeneric,
         target: TargetType = None,
-        additional_forward_args: object = None,
+        additional_forward_args: Optional[object] = None,
         feature_mask: Union[None, TensorOrTupleOfTensorsGeneric] = None,
         perturbations_per_eval: int = 1,
         show_progress: bool = False,
@@ -280,7 +280,7 @@ def attribute_future(
         self,
         inputs: TensorOrTupleOfTensorsGeneric,
         target: TargetType = None,
-        additional_forward_args: object = None,
+        additional_forward_args: Optional[object] = None,
         feature_mask: Union[None, TensorOrTupleOfTensorsGeneric] = None,
         perturbations_per_eval: int = 1,
         show_progress: bool = False,