code: #!/usr/bin/env python3
  
  import warnings
  from functools import partial
  from typing import Any, Callable, Dict, List, Optional, Tuple, Union
  
  import captum._utils.common as common
  import torch
  from captum._utils.av import AV
  from captum.attr import LayerActivation
  from captum.influence._core.influence import DataInfluence
  from torch import Tensor
  from torch.nn import Module
  from torch.utils.data import DataLoader, Dataset
  
  r"""
  Additional helper functions to calculate similarity metrics.
  """
  
  
  def euclidean_distance(test, train) -> 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, *).
      Returns pairwise euclidean distance Tensor of shape (batch_size_1, batch_size_2).
      """
      similarity = torch.cdist(
          test.view(test.shape[0], -1).unsqueeze(0),
          train.view(train.shape[0], -1).unsqueeze(0),
      ).squeeze(0)
      return similarity
  
  
  def cosine_similarity(test, train, replace_nan=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, *).
      Returns pairwise cosine similarity Tensor of shape (batch_size_1, batch_size_2).
      """
      test = test.view(test.shape[0], -1)
      train = train.view(train.shape[0], -1)
  
      if torch.__version__ <= "1.6.0":
          test_norm = torch.norm(test, p=None, dim=1, keepdim=True)
          train_norm = torch.norm(train, p=None, dim=1, keepdim=True)
      else:
          test_norm = torch.linalg.norm(test, ord=2, dim=1, keepdim=True)
          train_norm = torch.linalg.norm(train, ord=2, dim=1, keepdim=True)
  
      test = torch.where(test_norm != 0.0, test / test_norm, Tensor([replace_nan]))
      train = torch.where(train_norm != 0.0, train / train_norm, Tensor([replace_nan])).T
  
      similarity = torch.mm(test, train)
      return similarity
  
  
  r"""
  Implements abstract DataInfluence class and provides implementation details for
  similarity metric-based influence computation. Similarity metrics can be used to compare
  intermediate or final activation vectors of a model for different sets of input. Then,
  these can be used to draw conclusions about influential instances.
  
  Some standard similarity metrics such as dot product similarity or euclidean distance
  are provided, but the user can provide any custom similarity metric as well.
  """
  
  
  class SimilarityInfluence(DataInfluence):
      def __init__(
          self,
          module: Module,
          layers: Union[str, List[str]],
          influence_src_dataset: Dataset,
          activation_dir: str,
          model_id: str = "",
          similarity_metric: Callable = cosine_similarity,
          similarity_direction: str = "max",
          batch_size: int = 1,
          **kwargs: Any,
      ):
          r"""
          Args:
              module (torch.nn.Module): An instance of pytorch model. This model should
                      define all of its layers as attributes of the model.
              layers (str or List of str): The fully qualified layer(s) for which the
                      activation vectors are computed.
              influence_src_dataset (torch.utils.data.Dataset): PyTorch Dataset that is
                      used to create a PyTorch Dataloader to iterate over the dataset and
                      its labels. This is the dataset for which we will be seeking for
                      influential instances. In most cases this is the training dataset.
              activation_dir (str): The directory of the path to store
                      and retrieve activation computations. Best practice would be to use
                      an absolute path.
              model_id (str): The name/version of the model for which layer
                          activations are being computed. Activations will be stored and
                          loaded under the subdirectory with this name if provided.
              similarity_metric (Callable): This is a callable function that computes a
                      similarity metric between two representations. For example, the
                      representations pair could be from the training and test sets.
  
                      This function must adhere to certain standards. The inputs should be
                      torch Tensors with shape (batch_size_i/j, feature dimensions). The
                      output Tensor should have shape (batch_size_i, batch_size_j) with
                      scalar values corresponding to the similarity metric used for each
                      pairwise combination from the two batches.
  
                      For example, suppose we use `batch_size_1 = 16` for iterating
                      through `influence_src_dataset`, and for the `inputs` argument
                      we pass in a Tensor with 3 examples, i.e. batch_size_2 = 3. Also,
                      suppose that our inputs and intermediate activations throughout the
                      model will have dimension (N, C, H, W). Then, the feature dimensions
                      should be flattened within this function. For example::
  
                          >>> av_test.shape
                          torch.Size([3, N, C, H, W])
                          >>> av_src.shape
                          torch.Size([16, N, C, H, W])
                          >>> av_test = torch.view(av_test.shape[0], -1)
                          >>> av_test.shape
                          torch.Size([3, N x C x H x W])
  
                      and similarly for av_src. The similarity_metric should then use
                      these flattened tensors to return the pairwise similarity matrix.
                      For example, `similarity_metric(av_test, av_src)` should return a
                      tensor of shape (3, 16).
  
              batch_size (int): Batch size for iterating through `influence_src_dataset`.
              **kwargs: Additional key-value arguments that are necessary for specific
                      implementation of `DataInfluence` abstract class.
          """
          self.module = module
          self.layers = [layers] if isinstance(layers, str) else layers
          self.influence_src_dataset = influence_src_dataset
          self.activation_dir = activation_dir
          self.model_id = model_id
          self.batch_size = batch_size
  
          if similarity_direction == "max" or similarity_direction == "min":
              self.similarity_direction = similarity_direction
          else:
              raise ValueError(
                  f"{similarity_direction} is not a valid value. "
                  "Must be either 'max' or 'min'"
              )
  
          if similarity_metric is cosine_similarity:
              if "replace_nan" in kwargs:
                  self.replace_nan = kwargs["replace_nan"]
              else:
                  self.replace_nan = -2 if self.similarity_direction == "max" else 2
              similarity_metric = partial(cosine_similarity, replace_nan=self.replace_nan)
  
          self.similarity_metric = similarity_metric
  
          self.influence_src_dataloader = DataLoader(
              influence_src_dataset, batch_size, shuffle=False
          )
  
      def influence(  # type: ignore[override]
          self,
          inputs: Union[Tensor, Tuple[Tensor, ...]],
          top_k: int = 1,
          additional_forward_args: Optional[Any] = None,
          load_src_from_disk: bool = True,
          **kwargs: Any,
      ) -> Dict:
          r"""
          Args:
              inputs (tensor or tuple of tensors): Batch of examples for which influential
                      instances are computed. They are passed to the forward_func. The
                      first dimension in `inputs` tensor or tuple of tensors corresponds
                      to the batch size. A tuple of tensors is only passed in if this
                      is the input form that `module` accepts.
              top_k (int): The number of top-matching activations to return
              additional_forward_args (optional):  Additional arguments that will be
                      passed to forward_func after inputs.
              load_src_from_disk (bool): Loads activations for `influence_src_dataset`
                      where possible. Setting to False would force regeneration of
                      activations.
              load_input_from_disk (bool): Regenerates activations for inputs by default
                      and removes previous `inputs` activations that are flagged with
                      `inputs_id`. Setting to True will load prior matching inputs
                      activations. Note that this could lead to unexpected behavior if
                      `inputs_id` is not configured properly and activations are loaded
                      for a different, prior `inputs`.
              inputs_id (str): Used to identify inputs for loading activations.
  
              **kwargs: Additional key-value arguments that are necessary for specific
                      implementation of `DataInfluence` abstract class.
  
          Returns:
  
              influences (dict): Returns the influential instances retrieved from
              `influence_src_dataset` for each test example represented through a
              tensor or a tuple of tensor in `inputs`. Returned influential
              examples are represented as dict, with keys corresponding to
              the layer names passed in `layers`. Each value in the dict is a
              tuple containing the indices and values for the top k similarities
              from `influence_src_dataset` by the chosen metric. The first value
              in the tuple corresponds to the indices corresponding to the top k
              most similar examples, and the second value is the similarity score.
              The batch dimension corresponds to the batch dimension of `inputs`.
              If inputs.shape[0] == 5, then dict[`layer_name`][0].shape[0] == 5.
              These tensors will be of shape (inputs.shape[0], top_k).
          """
          inputs_batch_size = (
              inputs[0].shape[0] if isinstance(inputs, tuple) else inputs.shape[0]
          )
  
          influences: Dict[str, Any] = {}
  
          layer_AVDatasets = AV.generate_dataset_activations(
              self.activation_dir,
              self.module,
              self.model_id,
              self.layers,
              DataLoader(self.influence_src_dataset, self.batch_size, shuffle=False),
              identifier="src",
              load_from_disk=load_src_from_disk,
              return_activations=True,
          )
  
          assert layer_AVDatasets is not None and not isinstance(
              layer_AVDatasets, AV.AVDataset
          )
  
          layer_modules = [
              common._get_module_from_name(self.module, layer) for layer in self.layers
          ]
          test_activations = LayerActivation(self.module, layer_modules).attribute(
              inputs, additional_forward_args
          )
  
          minmax = self.similarity_direction == "max"
  
          # av_inputs shape: (inputs_batch_size, *) e.g. (inputs_batch_size, N, C, H, W)
          # av_src shape: (self.batch_size, *) e.g. (self.batch_size, N, C, H, W)
          test_activations = (
              test_activations if len(self.layers) > 1 else [test_activations]
          )
          for i, (layer, layer_AVDataset) in enumerate(
              zip(self.layers, layer_AVDatasets)
          ):
              topk_val, topk_idx = torch.Tensor(), torch.Tensor().long()
              zero_acts = torch.Tensor().long()
  
              av_inputs = test_activations[i]
              src_loader = DataLoader(layer_AVDataset)
              for j, av_src in enumerate(src_loader):
                  av_src = av_src.squeeze(0)
  
                  similarity = self.similarity_metric(av_inputs, av_src)
                  msg = (
                      "Output of custom similarity does not meet required dimensions. "
                      f"Your output has shape {similarity.shape}.\nPlease ensure the "
                      "output shape matches (inputs_batch_size, src_dataset_batch_size), "
                      f"which should be {(inputs_batch_size, self.batch_size)}."
                  )
                  assert similarity.shape == (inputs_batch_size, av_src.shape[0]), msg
                  if hasattr(self, "replace_nan"):
                      idx = (similarity == self.replace_nan).nonzero()
                      zero_acts = torch.cat((zero_acts, idx))
  
                  r"""
                  TODO: For models that can have tuples as activations, we should
                  allow similarity metrics to accept tuples, support topk selection.
                  """
  
                  topk_batch = min(top_k, self.batch_size)
                  values, indices = torch.topk(
                      similarity, topk_batch, dim=1, largest=minmax
                  )
                  indices += int(j * self.batch_size)
  
                  topk_val = torch.cat((topk_val, values), dim=1)
                  topk_idx = torch.cat((topk_idx, indices), dim=1)
  
                  # can modify how often to sort for efficiency? minor
                  sort_idx = torch.argsort(topk_val, dim=1, descending=minmax)
                  topk_val = torch.gather(topk_val, 1, sort_idx[:, :top_k])
                  topk_idx = torch.gather(topk_idx, 1, sort_idx[:, :top_k])
  
              influences[layer] = (topk_idx, topk_val)
  
              if torch.numel(zero_acts != 0):
                  zero_warning = (
                      f"Layer {layer} has zero-vector activations for some inputs. This "
                      "may cause undefined behavior for cosine similarity. The indices "
                      "for the offending inputs will be included under the key "
                      f"'zero_acts-{layer}' in the output dictionary. Indices are "
                      "returned as a tensor with [inputs_idx, src_dataset_idx] pairs "
                      "which may have corrupted similarity scores."
                  )
                  warnings.warn(zero_warning, RuntimeWarning)
                  key = "-".join(["zero_acts", layer])
                  influences[key] = zero_acts
  
          return influences
  
apis: torch.utils.data.DataLoader torch.linalg.norm torch.argsort torch.mm torch.gather torch.norm torch.topk torch.numel torch.cat torch.Tensor