Type Annotation for influence (#1247)

Summary:

as titled.

Reviewed By: vivekmig

Differential Revision: D55035833

Author: cicichen01

captum/influence/_core/arnoldi_influence_function.py

@@ -157,7 +157,7 @@ def _parameter_distill(
     k: Optional[int],
     hessian_reg: float,
     hessian_inverse_tol: float,
-):
+) -> Tuple[Tensor, List[Tuple[Tensor, ...]]]:
     """
     This takes the output of `_parameter_arnoldi`, and extracts the top-k eigenvalues
     / eigenvectors of the matrix that `_parameter_arnoldi` found the Krylov subspace

captum/influence/_core/influence_function.py

@@ -596,7 +596,7 @@ def _get_dataset_embeddings_intermediate_quantities_influence_function(
     batch_embeddings_fn: Callable,
     inputs_dataset: DataLoader,
     aggregate: bool,
-):
+) -> Tensor:
     """
     given `batch_embeddings_fn`, which produces the embeddings for a given batch,
     returns either the embeddings for an entire dataset (if `aggregate` is false),

captum/influence/_core/similarity_influence.py

@@ -18,7 +18,7 @@
 """
 
 
-def euclidean_distance(test, train) -> Tensor:
+def euclidean_distance(test: Tensor, train: Tensor) -> Tensor:
     r"""
     Calculates the pairwise euclidean distance for batches of feature vectors.
     Tensors test and train have shape (batch_size_1, *), and (batch_size_2, *).
@@ -31,7 +31,7 @@ def euclidean_distance(test, train) -> Tensor:
     return similarity
 
 
-def cosine_similarity(test, train, replace_nan=0) -> Tensor:
+def cosine_similarity(test: Tensor, train: Tensor, replace_nan: int = 0) -> Tensor:
     r"""
     Calculates the pairwise cosine similarity for batches of feature vectors.
     Tensors test and train have shape (batch_size_1, *), and (batch_size_2, *).

captum/influence/_core/tracincp_fast_rand_proj.py

@@ -720,7 +720,7 @@ def _basic_computation_tracincp_fast(
     targets: Tensor,
     loss_fn: Optional[Union[Module, Callable]] = None,
     reduction_type: Optional[str] = None,
-):
+) -> Tuple[Tensor, Tensor]:
     """
     For instances of TracInCPFast and children classes, computation of influence scores
     or self influence scores repeatedly calls this function for different checkpoints
@@ -1363,7 +1363,7 @@ def _set_projections_tracincp_fast_rand_proj(
     def _process_src_intermediate_quantities_tracincp_fast_rand_proj(
         self,
         src_intermediate_quantities: torch.Tensor,
-    ):
+    ) -> None:
         """
         Assumes `self._get_intermediate_quantities_tracin_fast_rand_proj` returns
         vector representations for each example, and that influence between a

captum/influence/_utils/common.py

@@ -4,6 +4,7 @@
 from typing import (
     Any,
     Callable,
+    Dict,
     Iterable,
     List,
     NamedTuple,
@@ -613,7 +614,7 @@ def _influence_batch_intermediate_quantities_influence_function(
     influence_inst: "IntermediateQuantitiesInfluenceFunction",
     test_batch: Tuple[Any, ...],
     train_batch: Tuple[Any, ...],
-):
+) -> Tensor:
     """
     computes influence of a test batch on a train batch, for implementations of
     `IntermediateQuantitiesInfluenceFunction`
@@ -628,7 +629,7 @@ def _influence_helper_intermediate_quantities_influence_function(
     influence_inst: "IntermediateQuantitiesInfluenceFunction",
     inputs_dataset: Union[Tuple[Any, ...], DataLoader],
     show_progress: bool,
-):
+) -> Tensor:
     """
     Helper function that computes influence scores for implementations of
     `NaiveInfluenceFunction` which implement the `compute_intermediate_quantities`
@@ -666,7 +667,7 @@ def _self_influence_helper_intermediate_quantities_influence_function(
     influence_inst: "IntermediateQuantitiesInfluenceFunction",
     inputs_dataset: Optional[Union[Tuple[Any, ...], DataLoader]],
     show_progress: bool,
-):
+) -> Tensor:
     """
     Helper function that computes self-influence scores for implementations of
     `NaiveInfluenceFunction` which implement the `compute_intermediate_quantities`
@@ -983,14 +984,14 @@ def _compute_batch_loss_influence_function_base(
         raise Exception
 
 
-def _set_attr(obj, names, val):
+def _set_attr(obj, names, val) -> None:
     if len(names) == 1:
         setattr(obj, names[0], val)
     else:
         _set_attr(getattr(obj, names[0]), names[1:], val)
 
 
-def _del_attr(obj, names):
+def _del_attr(obj, names) -> None:
     if len(names) == 1:
         delattr(obj, names[0])
     else:
@@ -1006,7 +1007,7 @@ def _model_make_functional(model, param_names, params):
     return params
 
 
-def _model_reinsert_params(model, param_names, params, register=False):
+def _model_reinsert_params(model, param_names, params, register: bool = False) -> None:
     for param_name, param in zip(param_names, params):
         _set_attr(
             model,
@@ -1024,7 +1025,7 @@ def _custom_functional_call(model, d, features):
     return out
 
 
-def _functional_call(model, d, features):
+def _functional_call(model: Module, d: Dict[str, Tensor], features):
     """
     Makes a call to `model.forward`, which is treated as a function of the parameters
     in `d`, a dict from parameter name to parameter, instead of as a function of