Type Hints Completion (#295)

Summary:
This PR completes type hints for remaining methods, adding type hints to utility methods, and making type signatures for methods with the return_convergence_delta flag clearer using Literal overrides.
Pull Request resolved: #295

Reviewed By: NarineK

Differential Revision: D20182905

Pulled By: vivekmig

fbshipit-source-id: 6889322297eff193c5614dec5fd6972ea41b2afe

Author: vivekmig

captum/attr/_core/deep_lift.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 import typing
-from typing import Optional, Tuple, Union, Any, List, Callable, cast
+from typing import Tuple, Union, Any, List, Callable, cast
 
 import warnings
 import torch
@@ -13,6 +13,7 @@
 import numpy as np
 
 from .._utils.common import (
+    _is_tuple,
     _format_input,
     _format_baseline,
     _format_callable_baseline,
@@ -30,7 +31,12 @@
 )
 from .._utils.attribution import GradientAttribution
 from .._utils.gradient import apply_gradient_requirements, undo_gradient_requirements
-from .._utils.typing import TensorOrTupleOfTensors
+from .._utils.typing import (
+    TensorOrTupleOfTensorsGeneric,
+    Literal,
+    TargetType,
+    BaselineType,
+)
 
 
 # Check if module backward hook can safely be used for the module that produced
@@ -76,43 +82,39 @@ def __init__(self, model: Module) -> None:
     @typing.overload
     def attribute(
         self,
-        inputs: TensorOrTupleOfTensors,
-        baselines: Union[
-            Tensor, int, float, Tuple[Union[Tensor, int, float], ...]
-        ] = None,
-        target: Optional[
-            Union[int, Tuple[int, ...], Tensor, List[Tuple[int, ...]]]
-        ] = None,
+        inputs: TensorOrTupleOfTensorsGeneric,
+        baselines: BaselineType = None,
+        target: TargetType = None,
         additional_forward_args: Any = None,
-        custom_attribution_func: Callable[..., Tuple[Tensor, ...]] = None,
-    ) -> TensorOrTupleOfTensors:
+        return_convergence_delta: Literal[False] = False,
+        custom_attribution_func: Union[None, Callable[..., Tuple[Tensor, ...]]] = None,
+    ) -> TensorOrTupleOfTensorsGeneric:
         ...
 
     @typing.overload
     def attribute(
         self,
-        inputs: TensorOrTupleOfTensors,
-        baselines: Optional[
-            Union[Tensor, int, float, Tuple[Union[Tensor, int, float], ...]]
-        ] = None,
-        target: Optional[
-            Union[int, Tuple[int, ...], Tensor, List[Tuple[int, ...]]]
-        ] = None,
+        inputs: TensorOrTupleOfTensorsGeneric,
+        baselines: BaselineType = None,
+        target: TargetType = None,
         additional_forward_args: Any = None,
-        return_convergence_delta: bool = False,
-        custom_attribution_func: Callable[..., Tuple[Tensor, ...]] = None,
-    ) -> Union[TensorOrTupleOfTensors, Tuple[TensorOrTupleOfTensors, Tensor]]:
+        *,
+        return_convergence_delta: Literal[True],
+        custom_attribution_func: Union[None, Callable[..., Tuple[Tensor, ...]]] = None,
+    ) -> Tuple[TensorOrTupleOfTensorsGeneric, Tensor]:
         ...
 
-    def attribute(
+    def attribute(  # type: ignore
         self,
-        inputs,
-        baselines=None,
-        target=None,
-        additional_forward_args=None,
-        return_convergence_delta=False,
-        custom_attribution_func=None,
-    ):
+        inputs: TensorOrTupleOfTensorsGeneric,
+        baselines: BaselineType = None,
+        target: TargetType = None,
+        additional_forward_args: Any = None,
+        return_convergence_delta: bool = False,
+        custom_attribution_func: Union[None, Callable[..., Tuple[Tensor, ...]]] = None,
+    ) -> Union[
+        TensorOrTupleOfTensorsGeneric, Tuple[TensorOrTupleOfTensorsGeneric, Tensor]
+    ]:
         r""""
         Implements DeepLIFT algorithm based on the following paper:
         Learning Important Features Through Propagating Activation Differences,
@@ -278,7 +280,7 @@ def attribute(
 
         # Keeps track whether original input is a tuple or not before
         # converting it into a tuple.
-        is_inputs_tuple = isinstance(inputs, tuple)
+        is_inputs_tuple = _is_tuple(inputs)
 
         inputs = _format_input(inputs)
         baselines = _format_baseline(baselines, inputs)
@@ -341,10 +343,8 @@ def attribute(
     def _construct_forward_func(
         self,
         forward_func: Callable,
-        inputs: TensorOrTupleOfTensors,
-        target: Optional[
-            Union[int, Tuple[int, ...], Tensor, List[Tuple[int, ...]]]
-        ] = None,
+        inputs: Tuple,
+        target: TargetType = None,
         additional_forward_args: Any = None,
     ) -> Callable:
         def forward_fn():
@@ -533,39 +533,45 @@ def __init__(self, model: Module) -> None:
     @typing.overload  # type: ignore
     def attribute(
         self,
-        inputs: TensorOrTupleOfTensors,
-        baselines: Union[TensorOrTupleOfTensors, Callable[..., TensorOrTupleOfTensors]],
-        target: Optional[
-            Union[int, Tuple[int, ...], Tensor, List[Tuple[int, ...]]]
-        ] = None,
+        inputs: TensorOrTupleOfTensorsGeneric,
+        baselines: Union[
+            TensorOrTupleOfTensorsGeneric, Callable[..., TensorOrTupleOfTensorsGeneric]
+        ],
+        target: TargetType = None,
         additional_forward_args: Any = None,
-        custom_attribution_func: Callable[..., Tuple[Tensor, ...]] = None,
-    ) -> TensorOrTupleOfTensors:
+        return_convergence_delta: Literal[False] = False,
+        custom_attribution_func: Union[None, Callable[..., Tuple[Tensor, ...]]] = None,
+    ) -> TensorOrTupleOfTensorsGeneric:
         ...
 
     @typing.overload
     def attribute(
         self,
-        inputs: TensorOrTupleOfTensors,
-        baselines: Union[TensorOrTupleOfTensors, Callable[..., TensorOrTupleOfTensors]],
-        target: Optional[
-            Union[int, Tuple[int, ...], Tensor, List[Tuple[int, ...]]]
-        ] = None,
+        inputs: TensorOrTupleOfTensorsGeneric,
+        baselines: Union[
+            TensorOrTupleOfTensorsGeneric, Callable[..., TensorOrTupleOfTensorsGeneric]
+        ],
+        target: TargetType = None,
         additional_forward_args: Any = None,
-        return_convergence_delta: bool = False,
-        custom_attribution_func: Callable[..., Tuple[Tensor, ...]] = None,
-    ) -> Union[TensorOrTupleOfTensors, Tuple[TensorOrTupleOfTensors, Tensor]]:
+        *,
+        return_convergence_delta: Literal[True],
+        custom_attribution_func: Union[None, Callable[..., Tuple[Tensor, ...]]] = None,
+    ) -> Tuple[TensorOrTupleOfTensorsGeneric, Tensor]:
         ...
 
-    def attribute(
+    def attribute(  # type: ignore
         self,
-        inputs,
-        baselines,
-        target=None,
-        additional_forward_args=None,
-        return_convergence_delta=False,
-        custom_attribution_func=None,
-    ):
+        inputs: TensorOrTupleOfTensorsGeneric,
+        baselines: Union[
+            TensorOrTupleOfTensorsGeneric, Callable[..., TensorOrTupleOfTensorsGeneric]
+        ],
+        target: TargetType = None,
+        additional_forward_args: Any = None,
+        return_convergence_delta: bool = False,
+        custom_attribution_func: Union[None, Callable[..., Tuple[Tensor, ...]]] = None,
+    ) -> Union[
+        TensorOrTupleOfTensorsGeneric, Tuple[TensorOrTupleOfTensorsGeneric, Tensor]
+    ]:
         r"""
         Extends DeepLift algorithm and approximates SHAP values using Deeplift.
         For each input sample it computes DeepLift attribution with respect to
@@ -724,7 +730,7 @@ def attribute(
 
         # Keeps track whether original input is a tuple or not before
         # converting it into a tuple.
-        is_inputs_tuple = isinstance(inputs, tuple)
+        is_inputs_tuple = _is_tuple(inputs)
 
         inputs = _format_input(inputs)
 
@@ -745,14 +751,18 @@ def attribute(
             exp_base,
             target=exp_tgt,
             additional_forward_args=exp_addit_args,
-            return_convergence_delta=return_convergence_delta,
+            return_convergence_delta=cast(
+                Literal[True, False], return_convergence_delta
+            ),
             custom_attribution_func=custom_attribution_func,
         )
         if return_convergence_delta:
-            attributions, delta = attributions
+            attributions, delta = cast(Tuple[Tuple[Tensor, ...], Tensor], attributions)
 
         attributions = tuple(
-            self._compute_mean_across_baselines(inp_bsz, base_bsz, attribution)
+            self._compute_mean_across_baselines(
+                inp_bsz, base_bsz, cast(Tensor, attribution)
+            )
             for attribution in attributions
         )
 
@@ -765,9 +775,11 @@ def _expand_inputs_baselines_targets(
         self,
         baselines: Tuple[Tensor, ...],
         inputs: Tuple[Tensor, ...],
-        target: Optional[Union[int, Tuple[int, ...], Tensor, List[Tuple[int, ...]]]],
+        target: TargetType,
         additional_forward_args: Any,
-    ):
+    ) -> Tuple[
+        Tuple[Tensor, ...], Tuple[Tensor, ...], TargetType, Any,
+    ]:
         inp_bsz = inputs[0].shape[0]
         base_bsz = baselines[0].shape[0]
 

captum/attr/_core/feature_ablation.py

@@ -4,20 +4,21 @@
 
 from torch import Tensor, dtype
 
-from typing import Callable, List, Optional, Tuple, Union, Any, cast
+from typing import Any, Callable, Tuple, Union, cast
 
 from .._utils.common import (
     _find_output_mode_and_verify,
     _format_attributions,
     _format_input,
     _format_input_baseline,
+    _is_tuple,
     _run_forward,
     _expand_additional_forward_args,
     _expand_target,
     _format_additional_forward_args,
 )
 from .._utils.attribution import PerturbationAttribution
-from .._utils.typing import TensorOrTupleOfTensors
+from .._utils.typing import TensorOrTupleOfTensorsGeneric, TargetType, BaselineType
 
 
 class FeatureAblation(PerturbationAttribution):
@@ -32,20 +33,15 @@ def __init__(self, forward_func: Callable) -> None:
         self.use_weights = False
 
     def attribute(
-        # type:ignore
         self,
-        inputs: TensorOrTupleOfTensors,
-        baselines: Optional[
-            Union[Tensor, int, float, Tuple[Union[Tensor, int, float], ...]]
-        ] = None,
-        target: Optional[
-            Union[int, Tuple[int, ...], Tensor, List[Tuple[int, ...]]]
-        ] = None,
+        inputs: TensorOrTupleOfTensorsGeneric,
+        baselines: BaselineType = None,
+        target: TargetType = None,
         additional_forward_args: Any = None,
-        feature_mask: Optional[TensorOrTupleOfTensors] = None,
+        feature_mask: Union[None, Tensor, Tuple[Tensor, ...]] = None,
         perturbations_per_eval: int = 1,
         **kwargs: Any
-    ) -> TensorOrTupleOfTensors:
+    ) -> TensorOrTupleOfTensorsGeneric:
         r""""
         A perturbation based approach to computing attribution, involving
         replacing each input feature with a given baseline / reference, and
@@ -228,7 +224,7 @@ def attribute(
         """
         # Keeps track whether original input is a tuple or not before
         # converting it into a tuple.
-        is_inputs_tuple = isinstance(inputs, tuple)
+        is_inputs_tuple = _is_tuple(inputs)
         inputs, baselines = _format_input_baseline(inputs, baselines)
         additional_forward_args = _format_additional_forward_args(
             additional_forward_args

captum/attr/_core/feature_permutation.py

@@ -1,11 +1,11 @@
 #!/usr/bin/env python3
-from typing import Any, Callable, List, Optional, Tuple, Union
+from typing import Any, Callable, Tuple, Union
 
 import torch
 from torch import Tensor
 
 from .feature_ablation import FeatureAblation
-from .._utils.typing import TensorOrTupleOfTensors
+from .._utils.typing import TargetType, TensorOrTupleOfTensorsGeneric
 
 
 def _permute_feature(x: Tensor, feature_mask: Tensor) -> Tensor:
@@ -67,15 +67,13 @@ def __init__(self, forward_func: Callable, perm_func: Callable = _permute_featur
     # signature to the parent
     def attribute(  # type: ignore
         self,
-        inputs: TensorOrTupleOfTensors,
-        target: Optional[
-            Union[int, Tuple[int, ...], Tensor, List[Tuple[int, ...]]]
-        ] = None,
+        inputs: TensorOrTupleOfTensorsGeneric,
+        target: TargetType = None,
         additional_forward_args: Any = None,
-        feature_mask: Optional[TensorOrTupleOfTensors] = None,
-        ablations_per_eval: int = 1,
+        feature_mask: Union[None, TensorOrTupleOfTensorsGeneric] = None,
+        perturbations_per_eval: int = 1,
         **kwargs: Any,
-    ) -> TensorOrTupleOfTensors:
+    ) -> TensorOrTupleOfTensorsGeneric:
         r"""
         This function is almost equivalent to `FeatureAblation.attribute`. The
         main difference is the way ablated examples are generated. Specifically
@@ -154,17 +152,17 @@ def attribute(  # type: ignore
                             each scalar within a tensor as a separate feature, which
                             is permuted independently.
                             Default: None
-                ablations_per_eval (int, optional): Allows permutations (ablations)
+                perturbations_per_eval (int, optional): Allows permutations
                             of multiple features to be processed simultaneously
                             in one call to forward_fn.  Each forward pass will
-                            contain a maximum of ablations_per_eval * #examples
+                            contain a maximum of perturbations_per_eval * #examples
                             samples.  For DataParallel models, each batch is
                             split among the available devices, so evaluations on
                             each available device contain at most
-                            (ablations_per_eval * #examples) / num_devices
+                            (perturbations_per_eval * #examples) / num_devices
                             samples.
                             If the forward function returns a single scalar per batch,
-                            ablations_per_eval must be set to 1.
+                            perturbations_per_eval must be set to 1.
                             Default: 1
                 **kwargs (Any, optional): Any additional arguments used by child
                             classes of FeatureAblation (such as Occlusion) to construct
@@ -179,7 +177,7 @@ def attribute(  # type: ignore
             target=target,
             additional_forward_args=additional_forward_args,
             feature_mask=feature_mask,
-            ablations_per_eval=ablations_per_eval,
+            perturbations_per_eval=perturbations_per_eval,
             **kwargs,
         )