Starter Task 2: Adding return type annotation to the _attribute_progress_setup method (#1576)

Summary:

The problem was that the _attribute_progress_setup method in feature_ablation.py didn't have a return type annotation. After making the relevant changes in the task 1 diff to only have the progress function return a tqdm object, I added a tqdm return type annotation to the _attribute_progress_setup function.

Differential Revision: D75972827

Author: ArjunR2404

captum/attr/_core/feature_ablation.py

@@ -30,7 +30,7 @@
     _run_forward,
 )
 from captum._utils.exceptions import FeatureAblationFutureError
-from captum._utils.progress import progress, SimpleProgress
+from captum._utils.progress import progress
 from captum._utils.typing import BaselineType, TargetType, TensorOrTupleOfTensorsGeneric
 from captum.attr._utils.attribution import PerturbationAttribution
 from captum.attr._utils.common import (
@@ -41,10 +41,7 @@
 from torch import dtype, Tensor
 from torch.futures import collect_all, Future
 
-try:
-    from tqdm.auto import tqdm
-except ImportError:
-    tqdm = None
+from tqdm.auto import tqdm
 
 IterableType = TypeVar("IterableType")
 
@@ -418,7 +415,7 @@ def _attribute_with_independent_feature_masks(
         formatted_feature_mask: Tuple[Tensor, ...],
         num_examples: int,
         perturbations_per_eval: int,
-        attr_progress: Optional[Union[SimpleProgress[IterableType], tqdm]],
+        attr_progress: Optional[tqdm],
         initial_eval: Tensor,
         flattened_initial_eval: Tensor,
         n_outputs: int,
@@ -500,7 +497,7 @@ def _attribute_with_cross_tensor_feature_masks(
         target: TargetType,
         baselines: BaselineType,
         formatted_feature_mask: Tuple[Tensor, ...],
-        attr_progress: Optional[Union[SimpleProgress[IterableType], tqdm]],
+        attr_progress: Optional[tqdm],
         flattened_initial_eval: Tensor,
         initial_eval: Tensor,
         n_outputs: int,
@@ -831,7 +828,7 @@ def _attribute_with_independent_feature_masks_future(
         baselines: BaselineType,
         formatted_feature_mask: Tuple[Tensor, ...],
         perturbations_per_eval: int,
-        attr_progress: Optional[Union[SimpleProgress[IterableType], tqdm]],
+        attr_progress: Optional[tqdm],
         processed_initial_eval_fut: Future[
             Tuple[List[Tensor], List[Tensor], Tensor, Tensor, int, dtype]
         ],
@@ -942,7 +939,7 @@ def _attribute_with_cross_tensor_feature_masks_future(
         target: TargetType,
         baselines: BaselineType,
         formatted_feature_mask: Tuple[Tensor, ...],
-        attr_progress: Optional[Union[SimpleProgress[IterableType], tqdm]],
+        attr_progress: Optional[tqdm],
         processed_initial_eval_fut: Future[
             Tuple[List[Tensor], List[Tensor], Tensor, Tensor, int, dtype]
         ],
@@ -1105,15 +1102,14 @@ def _fut_tuple_to_accumulate_fut_list_cross_tensor(
                 "_fut_tuple_to_accumulate_fut_list_cross_tensor failed"
             ) from e
 
-    # pyre-fixme[3] return type must be annotated
     def _attribute_progress_setup(
         self,
         formatted_inputs: Tuple[Tensor, ...],
         feature_mask: Tuple[Tensor, ...],
         enable_cross_tensor_attribution: bool,
         perturbations_per_eval: int,
         **kwargs: Any,
-    ):
+    ) -> tqdm:
         feature_counts = self._get_feature_counts(
             formatted_inputs, feature_mask, **kwargs
         )