class_docstr: A version of TracInCPFast which is optimized for "interactive" calls to
  `influence` for the purpose of calculating proponents / opponents, or
  influence scores. "Interactive" means there will be multiple calls to
  `influence`, with each call for a different batch of test examples, and
  subsequent calls rely on the results of previous calls. The implementation in
  this class has been optimized so that each call to `influence` is fast, so that
  it can be used for interactive analysis. This class should only be used for
  interactive use cases. It should not be used if `influence` will only be
  called once, because to enable fast calls to `influence`, time and memory
  intensive preprocessing is required in `__init__`. Furthermore, it should not
  be used to calculate self influence scores - `TracInCPFast` should be used
  instead for that purpose. To enable interactive analysis, this implementation
  computes and saves "embedding" vectors for all training examples in
  `train_dataset`. Crucially, the influence score of a training
  example on a test example is simply the dot-product of their corresponding
  vectors, and proponents / opponents can be found by first storing vectors for
  training examples in a nearest-neighbor data structure, and then finding the
  nearest-neighbors for a test example in terms of dot-product (see appendix F
  of the TracIn paper). This class should only be used if calls to `influence`
  to obtain proponents / opponents or influence scores will be made in an
  "interactive" manner, and there is sufficient memory to store vectors for the
  entire `train_dataset`. This is because in order to enable interactive
  analysis, this implementation incures overhead in `__init__` to setup the
  nearest-neighbors data structure, which is both time and memory intensive, as
  vectors corresponding to all training examples needed to be stored. To reduce
  memory usage, this implementation enables random projections of those vectors.
  Note that the influence scores computed with random projections are less
  accurate, though correct in expectation.
  
  In more detail regarding the "embedding" vectors - the influence of a training
  example on a test example, when only considering gradients in the last
  fully-connected layer, the sum of the contribution from each checkpoint. The
  contribution from a given checkpoint is
  :math`(x^T x')(\nabla_y f(y)^T \nabla_{y'} f(y'))`, using the notation in the
  description of `TracInCPFast`. As is, this is not a dot-product of 2 vectors.
  However, we can rewrite that contribution as
  :math`(x \nabla_y f(y)^T) \dot (x' f(y')^T)`. Both terms in this
  product are 2D matrices, as they are outer products, and the "product" is actually
  a dot-product, treating both matrices as vectors. Therefore, for a given
  checkpoint, its contribution to the "embedding" of an example is just the
  outer-product :math`(x \nabla_y f(y)^T)`, flattened. Furthemore, to reduce the
  dimension of this contribution, we can right-multiply and
  left-multiply the outer-product with two separate projection matrices. These
  transform :math`\nabla_y f(y)` and :math`x` to lower dimensional vectors. While
  the dimension of these two lower dimensional vectors do not necessarily need to
  be the same, in our implementation, we let them be the same, both equal to the
  square root of the desired projection dimension. Finally, the embedding of an
  example is the concatenation of the contributions from each checkpoint.
class_docstr_tok: ["A","version","of","TracInCPFast","which","is","optimized","for","``","interactive","''","calls","to","`","influence","`","for","the","purpose","of","calculating","proponents","\/","opponents",",","or","influence","scores",".","``","Interactive","''","means","there","will","be","multiple","calls","to","`","influence","`",",","with","each","call","for","a","different","batch","of","test","examples",",","and","subsequent","calls","rely","on","the","results","of","previous","calls",".","The","implementation","in","this","class","has","been","optimized","so","that","each","call","to","`","influence","`","is","fast",",","so","that","it","can","be","used","for","interactive","analysis",".","This","class","should","only","be","used","for","interactive","use","cases",".","It","should","not","be","used","if","`","influence","`","will","only","be","called","once",",","because","to","enable","fast","calls","to","`","influence","`",",","time","and","memory","intensive","preprocessing","is","required","in","`","__init__","`",".","Furthermore",",","it","should","not","be","used","to","calculate","self","influence","scores","-","`","TracInCPFast","`","should","be","used","instead","for","that","purpose",".","To","enable","interactive","analysis",",","this","implementation","computes","and","saves","``","embedding","''","vectors","for","all","training","examples","in","`","train_dataset","`",".","Crucially",",","the","influence","score","of","a","training","example","on","a","test","example","is","simply","the","dot-product","of","their","corresponding","vectors",",","and","proponents","\/","opponents","can","be","found","by","first","storing","vectors","for","training","examples","in","a","nearest-neighbor","data","structure",",","and","then","finding","the","nearest-neighbors","for","a","test","example","in","terms","of","dot-product","(","see","appendix","F","of","the","TracIn","paper",")",".","This","class","should","only","be","used","if","calls","to","`","influence","`","to","obtain","proponents","\/","opponents","or","influence","scores","will","be","made","in","an","''","interactive","''","manner",",","and","there","is","sufficient","memory","to","store","vectors","for","the","entire","`","train_dataset","`",".","This","is","because","in","order","to","enable","interactive","analysis",",","this","implementation","incures","overhead","in","`","__init__","`","to","setup","the","nearest-neighbors","data","structure",",","which","is","both","time","and","memory","intensive",",","as","vectors","corresponding","to","all","training","examples","needed","to","be","stored",".","To","reduce","memory","usage",",","this","implementation","enables","random","projections","of","those","vectors",".","Note","that","the","influence","scores","computed","with","random","projections","are","less","accurate",",","though","correct","in","expectation",".","In","more","detail","regarding","the","``","embedding","''","vectors","-","the","influence","of","a","training","example","on","a","test","example",",","when","only","considering","gradients","in","the","last","fully-connected","layer",",","the","sum","of","the","contribution","from","each","checkpoint",".","The","contribution","from","a","given","checkpoint","is",":","math","`","(","x^T","x","'",")","(","\\nabla_y","f","(","y",")","^T","\\nabla_","{","y","'","}","f","(","y","'",")",")","`",",","using","the","notation","in","the","description","of","`","TracInCPFast","`",".","As","is",",","this","is","not","a","dot-product","of","2","vectors",".","However",",","we","can","rewrite","that","contribution","as",":","math","`","(","x","\\nabla_y","f","(","y",")","^T",")","\\dot","(","x","'","f","(","y","'",")","^T",")","`",".","Both","terms","in","this","product","are","2D","matrices",",","as","they","are","outer","products",",","and","the","``","product","''","is","actually","a","dot-product",",","treating","both","matrices","as","vectors",".","Therefore",",","for","a","given","checkpoint",",","its","contribution","to","the","``","embedding","''","of","an","example","is","just","the","outer-product",":","math","`","(","x","\\nabla_y","f","(","y",")","^T",")","`",",","flattened",".","Furthemore",",","to","reduce","the","dimension","of","this","contribution",",","we","can","right-multiply","and","left-multiply","the","outer-product","with","two","separate","projection","matrices",".","These","transform",":","math","`","\\nabla_y","f","(","y",")","`","and",":","math","`","x","`","to","lower","dimensional","vectors",".","While","the","dimension","of","these","two","lower","dimensional","vectors","do","not","necessarily","need","to","be","the","same",",","in","our","implementation",",","we","let","them","be","the","same",",","both","equal","to","the","square","root","of","the","desired","projection","dimension",".","Finally",",","the","embedding","of","an","example","is","the","concatenation","of","the","contributions","from","each","checkpoint","."]
docstr: returns the variables `jacobian_projection` and `layer_input_projection`
  if needed, based on `self.projection_dim`. The two variables are
  used by `self._get_intermediate_quantities_fast_rand_proj`. They are both None
  if projection is not needed, due to the intermediate quantities (see the
  `_get_intermediate_quantities_fast_rand_proj` method for details) being no
  greater than `self.projection_dim` * C even without projection, where C is the
  number of checkpoints in the `checkpoints` argument to
  `TracInCPFastRandProj.__init__`.
  
  Args:
  
      dataloader (DataLoader): determining the projection requires knowing the
              dimensionality of the last layer's parameters (`jacobian_dim`
              below) and its input (`layer_input_dim` below). These are
              determined by passing a batch to `model`. `dataloader`
              provides that batch.
  
  Returns:
      jacobian_projection (Tensor or None): Projection matrix to apply to
              Jacobian of last layer to reduce its dimension, if needed.
              None otherwise.
      input_projection (Tensor or None): Projection matrix to apply to input of
              last layer to reduce its dimension, if needed. None otherwise.
docstr_tok: ["returns","the","variables","`","jacobian_projection","`","and","`","layer_input_projection","`","if","needed",",","based","on","`","self.projection_dim","`",".","The","two","variables","are","used","by","`","self._get_intermediate_quantities_fast_rand_proj","`",".","They","are","both","None","if","projection","is","not","needed",",","due","to","the","intermediate","quantities","(","see","the","`","_get_intermediate_quantities_fast_rand_proj","`","method","for","details",")","being","no","greater","than","`","self.projection_dim","`","*","C","even","without","projection",",","where","C","is","the","number","of","checkpoints","in","the","`","checkpoints","`","argument","to","`","TracInCPFastRandProj.__init__","`",".","Args",":","dataloader","(","DataLoader",")",":","determining","the","projection","requires","knowing","the","dimensionality","of","the","last","layer","'s","parameters","(","`","jacobian_dim","`","below",")","and","its","input","(","`","layer_input_dim","`","below",")",".","These","are","determined","by","passing","a","batch","to","`","model","`",".","`","dataloader","`","provides","that","batch",".","Returns",":","jacobian_projection","(","Tensor","or","None",")",":","Projection","matrix","to","apply","to","Jacobian","of","last","layer","to","reduce","its","dimension",",","if","needed",".","None","otherwise",".","input_projection","(","Tensor","or","None",")",":","Projection","matrix","to","apply","to","input","of","last","layer","to","reduce","its","dimension",",","if","needed",".","None","otherwise","."]
code: def _set_projections_tracincp_fast_rand_proj(
      self,
      dataloader: DataLoader,
  ) -> Optional[Tuple[torch.Tensor, torch.Tensor]]:
      """
      returns the variables `jacobian_projection` and `layer_input_projection`
      if needed, based on `self.projection_dim`. The two variables are
      used by `self._get_intermediate_quantities_fast_rand_proj`. They are both None
      if projection is not needed, due to the intermediate quantities (see the
      `_get_intermediate_quantities_fast_rand_proj` method for details) being no
      greater than `self.projection_dim` * C even without projection, where C is the
      number of checkpoints in the `checkpoints` argument to
      `TracInCPFastRandProj.__init__`.
  
      Args:
  
          dataloader (DataLoader): determining the projection requires knowing the
                  dimensionality of the last layer's parameters (`jacobian_dim`
                  below) and its input (`layer_input_dim` below). These are
                  determined by passing a batch to `model`. `dataloader`
                  provides that batch.
  
      Returns:
          jacobian_projection (Tensor or None): Projection matrix to apply to
                  Jacobian of last layer to reduce its dimension, if needed.
                  None otherwise.
          input_projection (Tensor or None): Projection matrix to apply to input of
                  last layer to reduce its dimension, if needed. None otherwise.
      """
      # figure out projection dimensions, if needed
  
      projection_dim = self.projection_dim
  
      projection_quantities = None
  
      if not (projection_dim is None):
          # figure out original dimensions by looking at data, passing through network
          self.checkpoints_load_func(
              self.model, next(iter(self.checkpoints))
          )
  
          batch = next(iter(dataloader))
          (
              batch_jacobians,
              batch_layer_inputs,
          ) = _basic_computation_tracincp_fast(
              self,
              batch[0:-1],
              batch[-1],
              self.loss_fn,
              self.reduction_type,
          )
  
          jacobian_dim = batch_jacobians.shape[
              1
          ]  # this is the dimension of the output of the last fully-connected layer
          layer_input_dim = batch_layer_inputs.shape[
              1
          ]  # this is the dimension of the input of the last fully-connected layer
          device = batch_jacobians.device
          dtype = batch_jacobians.dtype
  
          # choose projection if needed
          # without projection, the dimension of the intermediate quantities returned
          # by `_get_intermediate_quantities_fast_rand_proj` will be
          # `jacobian_dim` * `layer_input_dim` * number of checkpoints
          # this is because for each checkpoint, we compute a "partial" intermediate
          # quantity, and the intermediate quantity is the concatenation of the
          # "partial" intermediate quantities, and the dimension of each "partial"
          # intermediate quantity, without projection, is `jacobian_dim` *
          # `layer_input_dim`. However, `projection_dim` refers to the maximum
          # allowable dimension of the "partial" intermediate quantity. Therefore,
          # we only project if `jacobian_dim` * `layer_input_dim` > `projection_dim`.
          # `projection_dim` corresponds to the variable d in the top of page 15 of
          # the TracIn paper: https://arxiv.org/abs/2002.08484.
          if jacobian_dim * layer_input_dim > projection_dim:
              jacobian_projection_dim = min(
                  int(projection_dim**0.5), jacobian_dim
              )
              layer_input_projection_dim = min(
                  int(projection_dim**0.5), layer_input_dim
              )
              jacobian_projection = torch.normal(
                  torch.zeros(jacobian_dim, jacobian_projection_dim),
                  1.0 / jacobian_projection_dim**0.5,
              )
              layer_input_projection = torch.normal(
                  torch.zeros(
                      layer_input_dim, layer_input_projection_dim
                  ),
                  1.0 / layer_input_projection_dim**0.5,
              )
  
              projection_quantities = jacobian_projection.to(
                  device=device, dtype=dtype
              ), layer_input_projection.to(device=device, dtype=dtype)
  
      return projection_quantities
  
imports: import threading
  import warnings
  from collections import defaultdict
  from typing import Any, Callable, cast, Dict, Iterator, List, Optional, Tuple, Union
  import torch
  from captum._utils.common import _get_module_from_name, _sort_key_list
  from captum._utils.gradient import _gather_distributed_tensors
  from captum._utils.progress import NullProgress, progress
  from captum.influence._core.tracincp import _influence_route_to_helpers, KMostInfluentialResults, TracInCPBase
  from captum.influence._utils.common import _check_loss_fn, _format_inputs_dataset, _get_k_most_influential_helper, _jacobian_loss_wrt_inputs, _load_flexible_state_dict, _self_influence_by_batches_helper, _tensor_batch_dot
  from captum.influence._utils.nearest_neighbors import AnnoyNearestNeighbors, NearestNeighbors
  from captum.log import log_usage
  from torch import device, Tensor
  from torch.nn import Module
  from torch.utils.data import DataLoader, Dataset
label_desc: Use image node_id 6 for calling the TracInCPFastRandProj obj's underlying member method code with example usage: obj._set_projections_tracincp_fast_rand_proj(dataloader) and returns: projection_quantities