Enable pyre for Captum open source part- 1/2 (#1318)

Summary: Pull Request resolved: #1318

Differential Revision: D60748396

Author: yucu

captum/concept/__init__.py

@@ -1,4 +1,6 @@
 #!/usr/bin/env python3
+
+# pyre-strict
 from captum.concept._core.cav import CAV
 from captum.concept._core.concept import Concept, ConceptInterpreter
 from captum.concept._core.tcav import TCAV

captum/concept/_core/cav.py

@@ -1,5 +1,7 @@
 #!/usr/bin/env python3
 
+# pyre-strict
+
 import os
 from typing import Any, Dict, List, Optional
 
@@ -87,6 +89,7 @@ def assemble_save_path(
         file_name = concepts_to_str(concepts) + "-" + layer + ".pkl"
         return os.path.join(path, model_id, file_name)
 
+    # pyre-fixme[3]: Return type must be annotated.
     def save(self):
         r"""
         Saves a dictionary of the CAV computed values into a pickle file in the
@@ -134,6 +137,7 @@ def create_cav_dir_if_missing(save_path: str, model_id: str) -> None:
             os.makedirs(cav_model_id_path)
 
     @staticmethod
+    # pyre-fixme[3]: Return type must be annotated.
     def load(cavs_path: str, model_id: str, concepts: List[Concept], layer: str):
         r"""
         Loads CAV dictionary from a pickle file for given input

captum/concept/_core/concept.py

@@ -1,5 +1,7 @@
 #!/usr/bin/env python3
 
+# pyre-strict
+
 from typing import Callable, Union
 
 import torch
@@ -56,6 +58,7 @@ def __repr__(self) -> str:
         return "Concept(%r, %r)" % (self.id, self.name)
 
 
+# pyre-fixme[13]: Attribute `interpret` is never initialized.
 class ConceptInterpreter:
     r"""
     An abstract class that exposes an abstract interpret method
@@ -70,6 +73,7 @@ def __init__(self, model: Module) -> None:
         """
         self.model = model
 
+    # pyre-fixme[24]: Generic type `Callable` expects 2 type parameters.
     interpret: Callable
     r"""
     An abstract interpret method that performs concept-based model interpretability

captum/concept/_core/tcav.py

@@ -1,5 +1,7 @@
 #!/usr/bin/env python3
 
+# pyre-strict
+
 from collections import defaultdict
 from typing import Any, cast, Dict, List, Optional, Set, Tuple, Union
 
@@ -53,12 +55,17 @@ def __init__(self, datasets: List[AV.AVDataset], labels: List[int]) -> None:
         from itertools import accumulate
 
         offsets = [0] + list(accumulate(map(len, datasets), (lambda x, y: x + y)))
+        # pyre-fixme[4]: Attribute must be annotated.
         self.length = offsets[-1]
         self.datasets = datasets
         self.labels = labels
+        # pyre-fixme[4]: Attribute must be annotated.
         self.lowers = offsets[:-1]
+        # pyre-fixme[4]: Attribute must be annotated.
         self.uppers = offsets[1:]
 
+    # pyre-fixme[3]: Return type must be annotated.
+    # pyre-fixme[2]: Parameter must be annotated.
     def _i_to_k(self, i):
 
         left, right = 0, len(self.uppers)
@@ -71,6 +78,7 @@ def _i_to_k(self, i):
             else:
                 right = mid
 
+    # pyre-fixme[3]: Return type must be annotated.
     def __getitem__(self, i: int):
         """
         Returns a batch of activation vectors, as well as a batch of labels
@@ -89,8 +97,14 @@ def __getitem__(self, i: int):
         assert i < self.length
         k = self._i_to_k(i)
         inputs = self.datasets[k][i - self.lowers[k]]
+        # pyre-fixme[16]: Item `tuple` of `Union[Tensor, Tuple[Tensor, ...]]` has no
+        #  attribute `shape`.
         assert len(inputs.shape) == 2
 
+        # pyre-fixme[16]: Item `tuple` of `Union[Tensor, Tuple[Tensor, ...]]` has no
+        #  attribute `size`.
+        # pyre-fixme[16]: Item `tuple` of `Union[Tensor, Tuple[Tensor, ...]]` has no
+        #  attribute `device`.
         labels = torch.tensor([self.labels[k]] * inputs.size(0), device=inputs.device)
         return inputs, labels
 
@@ -102,11 +116,14 @@ def __len__(self) -> int:
 
 
 def train_cav(
+    # pyre-fixme[2]: Parameter must be annotated.
     model_id,
     concepts: List[Concept],
     layers: Union[str, List[str]],
     classifier: Classifier,
     save_path: str,
+    # pyre-fixme[24]: Generic type `dict` expects 2 type parameters, use
+    #  `typing.Dict[<key type>, <value type>]` to avoid runtime subscripting errors.
     classifier_kwargs: Dict,
 ) -> Dict[str, Dict[str, CAV]]:
     r"""
@@ -159,8 +176,11 @@ def train_cav(
 
         labels = [concept.id for concept in concepts]
 
+        # pyre-fixme[22]: The cast is redundant.
         labelled_dataset = LabelledDataset(cast(List[AV.AVDataset], datasets), labels)
 
+        # pyre-fixme[3]: Return type must be annotated.
+        # pyre-fixme[2]: Parameter must be annotated.
         def batch_collate(batch):
             inputs, labels = zip(*batch)
             return torch.cat(inputs), torch.cat(labels)
@@ -249,6 +269,7 @@ def __init__(
         model_id: str = "default_model_id",
         classifier: Optional[Classifier] = None,
         layer_attr_method: Optional[LayerAttribution] = None,
+        # pyre-fixme[2]: Parameter must be annotated.
         attribute_to_layer_input=False,
         save_path: str = "./cav/",
         **classifier_kwargs: Any,
@@ -300,19 +321,28 @@ def __init__(
             For more thorough examples, please check out TCAV tutorial and test cases.
         """
         ConceptInterpreter.__init__(self, model)
+        # pyre-fixme[4]: Attribute must be annotated.
         self.layers = [layers] if isinstance(layers, str) else layers
         self.model_id = model_id
         self.concepts: Set[Concept] = set()
         self.classifier = classifier
+        # pyre-fixme[4]: Attribute must be annotated.
         self.classifier_kwargs = classifier_kwargs
+        # pyre-fixme[8]: Attribute has type `Dict[str, Dict[str, CAV]]`; used as
+        #  `DefaultDict[Variable[_KT], DefaultDict[Variable[_KT], Variable[_VT]]]`.
         self.cavs: Dict[str, Dict[str, CAV]] = defaultdict(lambda: defaultdict())
         if self.classifier is None:
             self.classifier = DefaultClassifier()
         if layer_attr_method is None:
+            # pyre-fixme[4]: Attribute must be annotated.
             self.layer_attr_method = cast(
                 LayerAttribution,
                 LayerGradientXActivation(  # type: ignore
-                    model, None, multiply_by_inputs=False
+                    model,
+                    # pyre-fixme[6]: For 2nd argument expected `ModuleOrModuleList`
+                    #  but got `None`.
+                    None,
+                    multiply_by_inputs=False,
                 ),
             )
         else:
@@ -324,6 +354,7 @@ def __init__(
             "will use `default_model_id` as its default value."
         )
 
+        # pyre-fixme[4]: Attribute must be annotated.
         self.attribute_to_layer_input = attribute_to_layer_input
         self.save_path = save_path
 
@@ -341,6 +372,8 @@ def generate_all_activations(self) -> None:
         for concept in self.concepts:
             self.generate_activation(self.layers, concept)
 
+    # pyre-fixme[24]: Generic type `list` expects 1 type parameter, use
+    #  `typing.List[<element type>]` to avoid runtime subscripting errors.
     def generate_activation(self, layers: Union[str, List], concept: Concept) -> None:
         r"""
         Computes layer activations for the specified `concept` and
@@ -361,6 +394,8 @@ def generate_activation(self, layers: Union[str, List], concept: Concept) -> Non
             "Data iterator for concept id:",
             "{} must be specified".format(concept.id),
         )
+        # pyre-fixme[6]: For 1st argument expected `Iterable[Variable[_T]]` but got
+        #  `Optional[DataLoader[typing.Any]]`.
         for i, examples in enumerate(concept.data_iter):
             activations = layer_act.attribute.__wrapped__(  # type: ignore
                 layer_act,
@@ -447,6 +482,7 @@ def load_cavs(
                         concept_layers[concept].append(layer)
         return layers, concept_layers
 
+    # pyre-fixme[3]: Return type must be annotated.
     def compute_cavs(
         self,
         experimental_sets: List[List[Concept]],
@@ -566,6 +602,7 @@ def interpret(
         inputs: TensorOrTupleOfTensorsGeneric,
         experimental_sets: List[List[Concept]],
         target: TargetType = None,
+        # pyre-fixme[2]: Parameter annotation cannot be `Any`.
         additional_forward_args: Any = None,
         processes: Optional[int] = None,
         **kwargs: Any,
@@ -661,6 +698,9 @@ def interpret(
         )
         self.compute_cavs(experimental_sets, processes=processes)
 
+        # pyre-fixme[9]: scores has type `Dict[str, Dict[str, Dict[str, Tensor]]]`;
+        #  used as `DefaultDict[Variable[_KT], DefaultDict[Variable[_KT],
+        #  Variable[_VT]]]`.
         scores: Dict[str, Dict[str, Dict[str, Tensor]]] = defaultdict(
             lambda: defaultdict()
         )
@@ -704,6 +744,7 @@ def interpret(
             attribs = _format_tensor_into_tuples(attribs)
             # n_inputs x n_features
             attribs = torch.cat(
+                # pyre-fixme[16]: `None` has no attribute `__iter__`.
                 [torch.reshape(attrib, (attrib.shape[0], -1)) for attrib in attribs],
                 dim=1,
             )
@@ -713,6 +754,7 @@ def interpret(
             classes = []
             for concepts in experimental_sets:
                 concepts_key = concepts_to_str(concepts)
+                # pyre-fixme[33]: Given annotation cannot contain `Any`.
                 cavs_stats = cast(Dict[str, Any], self.cavs[concepts_key][layer].stats)
                 cavs.append(cavs_stats["weights"].float().detach().tolist())
                 classes.append(cavs_stats["classes"])

captum/concept/_utils/classifier.py

@@ -1,5 +1,7 @@
 #!/usr/bin/env python3
 
+# pyre-strict
+
 import random
 import warnings
 from abc import ABC, abstractmethod
@@ -64,7 +66,11 @@ def __init__(self) -> None:
 
     @abstractmethod
     def train_and_eval(
-        self, dataloader: DataLoader, **kwargs: Any
+        self,
+        dataloader: DataLoader,
+        **kwargs: Any
+        # pyre-fixme[24]: Generic type `dict` expects 2 type parameters, use
+        # `typing.Dict[<key type>, <value type>]` to avoid runtime subscripting errors.
     ) -> Union[Dict, None]:
         r"""
         This method is responsible for training a classifier using the data
@@ -132,12 +138,18 @@ def __init__(self) -> None:
             " both train and test datasets in the memory. Consider defining"
             " your own classifier that doesn't rely heavily on memory, for"
             " large number of concepts, by extending"
-            " `Classifer` abstract class"
+            " `Classifer` abstract class",
+            stacklevel=2,
         )
         self.lm = model.SkLearnSGDClassifier(alpha=0.01, max_iter=1000, tol=1e-3)
 
     def train_and_eval(
-        self, dataloader: DataLoader, test_split_ratio: float = 0.33, **kwargs: Any
+        self,
+        dataloader: DataLoader,
+        test_split_ratio: float = 0.33,
+        **kwargs: Any
+        # pyre-fixme[24]: Generic type `dict` expects 2 type parameters, use
+        # `typing.Dict[<key type>, <value type>]` to avoid runtime subscripting errors.
     ) -> Union[Dict, None]:
         r"""
          Implements Classifier::train_and_eval abstract method for small concept
@@ -169,6 +181,7 @@ def train_and_eval(
             inputs.append(input)
             labels.append(label)
 
+        # pyre-fixme[61]: `input` is undefined, or not always defined.
         device = "cpu" if input is None else input.device
         x_train, x_test, y_train, y_test = _train_test_split(
             torch.cat(inputs), torch.cat(labels), test_split=test_split_ratio

captum/concept/_utils/common.py

@@ -1,5 +1,7 @@
 #!/usr/bin/env python3
 
+# pyre-strict
+
 from typing import List
 
 from captum.concept._core.concept import Concept

captum/concept/_utils/data_iterator.py

@@ -1,5 +1,7 @@
 #!/usr/bin/env python3
 
+# pyre-strict
+
 import glob
 import os
 from typing import Callable, Iterator
@@ -13,6 +15,7 @@ class CustomIterableDataset(IterableDataset):
     An auxiliary class for iterating through a dataset.
     """
 
+    # pyre-fixme[24]: Generic type `Callable` expects 2 type parameters.
     def __init__(self, transform_filename_to_tensor: Callable, path: str) -> None:
         r"""
         Args:
@@ -21,6 +24,7 @@ def __init__(self, transform_filename_to_tensor: Callable, path: str) -> None:
             path (str): Path to dataset files. This can be either a path to a
                         directory or a file where input examples are stored.
         """
+        # pyre-fixme[4]: Attribute must be annotated.
         self.file_itr = None
         self.path = path
 

captum/influence/__init__.py

@@ -1,5 +1,7 @@
 #!/usr/bin/env python3
 
+# pyre-strict
+
 from captum.influence._core.influence import DataInfluence
 from captum.influence._core.influence_function import NaiveInfluenceFunction
 from captum.influence._core.similarity_influence import SimilarityInfluence