Fix layer deeplift pyre fixme issues

Differential Revision: D67705583

Author: Vivek Miglani

captum/attr/_core/layer/layer_deep_lift.py

@@ -2,7 +2,18 @@
 
 # pyre-strict
 import typing
-from typing import Any, Callable, cast, Dict, Literal, Optional, Sequence, Tuple, Union
+from typing import (
+    Any,
+    Callable,
+    cast,
+    Dict,
+    List,
+    Literal,
+    Optional,
+    Sequence,
+    Tuple,
+    Union,
+)
 
 import torch
 from captum._utils.common import (
@@ -96,6 +107,7 @@ def __init__(
 
     # Ignoring mypy error for inconsistent signature with DeepLift
     @typing.overload  # type: ignore
+    @log_usage()
     def attribute(
         self,
         inputs: Union[Tensor, Tuple[Tensor, ...]],
@@ -110,6 +122,7 @@ def attribute(
     ) -> Tuple[Union[Tensor, Tuple[Tensor, ...]], Tensor]: ...
 
     @typing.overload
+    @log_usage()
     def attribute(
         self,
         inputs: Union[Tensor, Tuple[Tensor, ...]],
@@ -123,8 +136,6 @@ def attribute(
     ) -> Union[Tensor, Tuple[Tensor, ...]]: ...
 
     @log_usage()
-    # pyre-fixme[43]: This definition does not have the same decorators as the
-    #  preceding overload(s).
     def attribute(
         self,
         inputs: Union[Tensor, Tuple[Tensor, ...]],
@@ -321,8 +332,9 @@ def attribute(
                 additional_forward_args,
             )
 
-            # pyre-fixme[24]: Generic type `Sequence` expects 1 type parameter.
-            def chunk_output_fn(out: TensorOrTupleOfTensorsGeneric) -> Sequence:
+            def chunk_output_fn(
+                out: TensorOrTupleOfTensorsGeneric,
+            ) -> Sequence[Union[Tensor, List[Tensor]]]:
                 if isinstance(out, Tensor):
                     return out.chunk(2)
                 return tuple(out_sub.chunk(2) for out_sub in out)
@@ -434,8 +446,7 @@ def __init__(
 
     # Ignoring mypy error for inconsistent signature with DeepLiftShap
     @typing.overload  # type: ignore
-    # pyre-fixme[43]: The implementation of `attribute` does not accept all possible
-    #  arguments of overload defined on line `453`.
+    @log_usage()
     def attribute(
         self,
         inputs: Union[Tensor, Tuple[Tensor, ...]],
@@ -451,8 +462,7 @@ def attribute(
     ) -> Tuple[Union[Tensor, Tuple[Tensor, ...]], Tensor]: ...
 
     @typing.overload
-    # pyre-fixme[43]: The implementation of `attribute` does not accept all possible
-    #  arguments of overload defined on line `439`.
+    @log_usage()
     def attribute(
         self,
         inputs: Union[Tensor, Tuple[Tensor, ...]],
@@ -467,8 +477,6 @@ def attribute(
     ) -> Union[Tensor, Tuple[Tensor, ...]]: ...
 
     @log_usage()
-    # pyre-fixme[43]: This definition does not have the same decorators as the
-    #  preceding overload(s).
     def attribute(
         self,
         inputs: Union[Tensor, Tuple[Tensor, ...]],
@@ -654,7 +662,7 @@ def attribute(
         ) = DeepLiftShap._expand_inputs_baselines_targets(
             self, baselines, inputs, target, additional_forward_args
         )
-        attributions = LayerDeepLift.attribute.__wrapped__(  # type: ignore
+        attribs_layer_deeplift = LayerDeepLift.attribute.__wrapped__(  # type: ignore
             self,
             exp_inp,
             exp_base,
@@ -667,8 +675,12 @@ def attribute(
             attribute_to_layer_input=attribute_to_layer_input,
             custom_attribution_func=custom_attribution_func,
         )
+        delta: Tensor
+        attributions: Union[Tensor, Tuple[Tensor, ...]]
         if return_convergence_delta:
-            attributions, delta = attributions
+            attributions, delta = attribs_layer_deeplift
+        else:
+            attributions = attribs_layer_deeplift
         if isinstance(attributions, tuple):
             attributions = tuple(
                 DeepLiftShap._compute_mean_across_baselines(
@@ -681,15 +693,17 @@ def attribute(
                 self, inp_bsz, base_bsz, attributions
             )
         if return_convergence_delta:
-            # pyre-fixme[61]: `delta` is undefined, or not always defined.
             return attributions, delta
         else:
-            # pyre-fixme[7]: Expected `Union[Tuple[Union[Tensor,
-            #  typing.Tuple[Tensor, ...]], Tensor], Tensor, typing.Tuple[Tensor, ...]]`
-            #  but got `Union[tuple[Tensor], Tensor]`.
-            return attributions
+            return cast(
+                Union[
+                    Tensor,
+                    Tuple[Tensor, ...],
+                    Tuple[Union[Tensor, Tuple[Tensor, ...]], Tensor],
+                ],
+                attributions,
+            )
 
     @property
-    # pyre-fixme[3]: Return type must be annotated.
     def multiplies_by_inputs(self) -> bool:
         return self._multiply_by_inputs