class_docstr: This is a computationally-efficient implementation that computes the type of
  "infinitesimal" influence scored defined in the paper "Understanding Black-box
  Predictions via Influence Functions" by Koh et al
  (https://arxiv.org/pdf/1703.04730.pdf). This implementation does *not* follow
  the approach in that paper, however. Instead, it follows an implementation that is
  several orders of magnitudes faster, described in the paper "Scaling Up Influence
  Functions" by Schioppa et al (https://arxiv.org/pdf/2112.03052.pdf).
  
  This implementation computes a low-rank approximation of the inverse Hessian, i.e.
  a tall and skinny (with width k) matrix :math`R` such that
  :math`H^{-1} \approx RR'`, where k is small. In particular, let :math`V` be the
  matrix of width k whose columns contain the top-k eigenvectors of :math`H`, and let
  :math`S` be the k by k matrix whose diagonals contain the corresponding eigenvalues.
  This implementation lets :math`R=VS^{-0.5}`. Thus, the core computational step is
  computing the top-k eigenvalues / eigenvectors.
  
  This approximation is useful for several reasons:
  - It avoids numerical issues associated with inverting small eigenvalues
  - Since the influence score is given by
    :math`\nabla_\theta L(x)' H^{-1} \nabla_\theta L(z)`, which is approximated by
    :math`(\nabla_\theta L(x)' R) (\nabla_\theta L(z)' R)`, we can compute an
    "influence embedding" for a given example :math`x`, :math`\nabla_\theta L(x)' R`,
    such that the influence score of one example on another is approximately the
    dot-product of their respective embeddings.
  - Even for large models, we can store `R` in memory, provided k is small. This
    means influence embeddings (and thus influence scores) can be efficiently
    computed by doing a backwards pass to compute :math`\nabla_\theta L(x)` and then
    multiplying by :math`R'`. This is orders of magnitude faster than the previous
    LISSA approach of Koh et al, which to compute the influence score involving a
    given example, need to compute Hessian-vector products involving on the order
    of 10^4 examples.
  
  The key novelty of the approach by Schioppa et al is that it uses the Arnoldi
  iteration to find the top-k eigenvalues / eigenvectors of the Hessian without
  explicitly forming the Hessian. In more detail, the approach first runs the
  Arnoldi iteration, which only requires the ability to compute Hessian-vector
  products, to find a Krylov subspace of moderate dimension, i.e. 200. It then finds
  the top-k eigevalues / eigenvectors of the restriction of the Hessian to the
  subspace, where k is small, i.e. 50. Finally, it expresses the eigenvectors in
  the original basis. This approach for finding the top-k eigenvalues / eigenvectors
  is justified by the property of the Arnoldi iteration, that the Krylov subspace
  it returns tends to contain the top eigenvectors.
  
  This implementation require some computation time `__init__`, where it
  runs the Arnoldi iteration to calculate `R`. This computation is linear in
  `arnoldi_dim` as well as the size of `hessian_dataset`. After that initial
  overhead, calculation of influence scores is quick, only requiring a backwards pass
  and multiplication, per example.
  
  Unlike `NaiveInfluenceFunction`, this implementation does not flatten any
  parameters, as the 2D Hessian is never formed, and Pytorch's Hessian-vector
  implementation (`torch.autograd.functional.hvp`) allows the input and output
  vector to be a tuple of tensors. Avoiding flattening / unflattening parameters
  brings scalability gains.
docstr: Computes "embedding" vectors for all examples in a single batch, or a
  `Dataloader` that yields batches. These embedding vectors are constructed so
  that the influence score of a training example on a test example is simply the
  dot-product of their corresponding vectors. In both cases, a batch should be
  small enough so that a backwards pass for a batch does not lead to
  out-of-memory errors.
  
  In more detail, the embedding vector for an example `x` is
  :math`\nabla_\theta L(x)' R`, where :math`R` is as defined in this class'
  description. Each element of `R` and :math`\nabla_\theta L(x)` lie in
  parameter-space. Therefore, if parameter-space were 1D, so that `R` were
  a 2D tensor whose columns are different elements in parameter-space, we would
  compute the embeddings for a batch by assembling :math`\nabla_\theta L(x)` for
  all examples `x` in the batch as rows in a 2D "batch parameter-gradient"
  tensor, and right-multiplying by `R`. However, parameter-space in this
  implementation is actually a tuple of tensors. So we do the analogous
  computation given this representation of parameter-space.
  
  If `aggregate` is True, the *sum* of the vectors for all examples is returned,
  instead of the vectors for each example. This can be useful for computing the
  influence of a given training example on the total loss over a validation
  dataset, because due to properties of the dot-product, this influence is the
  dot-product of the training example's vector with the sum of the vectors in the
  validation dataset. Also, by doing the sum aggregation within this method as
  opposed to outside of it (by computing all vectors for the validation dataset,
  then taking the sum) allows memory usage to be reduced.
  
  Args:
      inputs_dataset (Tuple, or DataLoader): Either a single tuple of any, or a
              `DataLoader`, where each batch yielded is a tuple of any. In
              either case, the tuple represents a single batch, where the last
              element is assumed to be the labels for the batch. That is,
              `model(*batch[0:-1])` produces the output for `model`, and
              and `batch[-1]` are the labels, if any. Here, `model` is model
              provided in initialization. This is the same assumption made for
              each batch yielded by training dataset `train_dataset`.
      aggregate (bool): Whether to return the sum of the vectors for all
              examples, as opposed to vectors for each example.
      show_progress (bool, optional): Computation of vectors can take a long
              time if `inputs_dataset` represents many examples. If
              `show_progress`is true, the progress of this computation will be
              displayed. In particular, the number of batches for which
              vectors have been computed will be displayed. It will try to
              use tqdm if available for advanced features (e.g. time estimation).
              Otherwise, it will fallback to a simple output of progress.
              Default: False
      return_on_cpu (bool, optional): Whether to return the vectors on the cpu
              (or if not, the gpu). If None, is set to the device that the model
              is on.
              Default: None
      test (bool, optional): Whether to compute the vectors using the loss
              function `test_loss_fn` provided in initialization (instead of
              `loss_fn`). This argument does not matter if `test_loss_fn` was
              not provided, as in this case, `test_loss_fn` and `loss_fn` are the
              same.
  
  Returns:
      intermediate_quantities (Tensor): This is a 2D tensor with shape
              `(N, projection_dim)`, where `N` is the total number of examples in
              `inputs_dataset`, and `projection_dim` was provided in
              initialization. Each row contains the vector for a different
              example.
docstr_tok: ["Computes","``","embedding","''","vectors","for","all","examples","in","a","single","batch",",","or","a","`","Dataloader","`","that","yields","batches",".","These","embedding","vectors","are","constructed","so","that","the","influence","score","of","a","training","example","on","a","test","example","is","simply","the","dot-product","of","their","corresponding","vectors",".","In","both","cases",",","a","batch","should","be","small","enough","so","that","a","backwards","pass","for","a","batch","does","not","lead","to","out-of-memory","errors",".","In","more","detail",",","the","embedding","vector","for","an","example","`","x","`","is",":","math","`","\\nabla_\\theta","L","(","x",")","'","R","`",",","where",":","math","`","R","`","is","as","defined","in","this","class'","description",".","Each","element","of","`","R","`","and",":","math","`","\\nabla_\\theta","L","(","x",")","`","lie","in","parameter-space",".","Therefore",",","if","parameter-space","were","1D",",","so","that","`","R","`","were","a","2D","tensor","whose","columns","are","different","elements","in","parameter-space",",","we","would","compute","the","embeddings","for","a","batch","by","assembling",":","math","`","\\nabla_\\theta","L","(","x",")","`","for","all","examples","`","x","`","in","the","batch","as","rows","in","a","2D","``","batch","parameter-gradient","''","tensor",",","and","right-multiplying","by","`","R","`",".","However",",","parameter-space","in","this","implementation","is","actually","a","tuple","of","tensors",".","So","we","do","the","analogous","computation","given","this","representation","of","parameter-space",".","If","`","aggregate","`","is","True",",","the","*","sum","*","of","the","vectors","for","all","examples","is","returned",",","instead","of","the","vectors","for","each","example",".","This","can","be","useful","for","computing","the","influence","of","a","given","training","example","on","the","total","loss","over","a","validation","dataset",",","because","due","to","properties","of","the","dot-product",",","this","influence","is","the","dot-product","of","the","training","example","'s","vector","with","the","sum","of","the","vectors","in","the","validation","dataset",".","Also",",","by","doing","the","sum","aggregation","within","this","method","as","opposed","to","outside","of","it","(","by","computing","all","vectors","for","the","validation","dataset",",","then","taking","the","sum",")","allows","memory","usage","to","be","reduced",".","Args",":","inputs_dataset","(","Tuple",",","or","DataLoader",")",":","Either","a","single","tuple","of","any",",","or","a","`","DataLoader","`",",","where","each","batch","yielded","is","a","tuple","of","any",".","In","either","case",",","the","tuple","represents","a","single","batch",",","where","the","last","element","is","assumed","to","be","the","labels","for","the","batch",".","That","is",",","`","model","(","*","batch","[","0",":","-1","]",")","`","produces","the","output","for","`","model","`",",","and","and","`","batch","[","-1","]","`","are","the","labels",",","if","any",".","Here",",","`","model","`","is","model","provided","in","initialization",".","This","is","the","same","assumption","made","for","each","batch","yielded","by","training","dataset","`","train_dataset","`",".","aggregate","(","bool",")",":","Whether","to","return","the","sum","of","the","vectors","for","all","examples",",","as","opposed","to","vectors","for","each","example",".","show_progress","(","bool",",","optional",")",":","Computation","of","vectors","can","take","a","long","time","if","`","inputs_dataset","`","represents","many","examples",".","If","`","show_progress","`","is","true",",","the","progress","of","this","computation","will","be","displayed",".","In","particular",",","the","number","of","batches","for","which","vectors","have","been","computed","will","be","displayed",".","It","will","try","to","use","tqdm","if","available","for","advanced","features","(","e.g",".","time","estimation",")",".","Otherwise",",","it","will","fallback","to","a","simple","output","of","progress",".","Default",":","False","return_on_cpu","(","bool",",","optional",")",":","Whether","to","return","the","vectors","on","the","cpu","(","or","if","not",",","the","gpu",")",".","If","None",",","is","set","to","the","device","that","the","model","is","on",".","Default",":","None","test","(","bool",",","optional",")",":","Whether","to","compute","the","vectors","using","the","loss","function","`","test_loss_fn","`","provided","in","initialization","(","instead","of","`","loss_fn","`",")",".","This","argument","does","not","matter","if","`","test_loss_fn","`","was","not","provided",",","as","in","this","case",",","`","test_loss_fn","`","and","`","loss_fn","`","are","the","same",".","Returns",":","intermediate_quantities","(","Tensor",")",":","This","is","a","2D","tensor","with","shape","`","(","N",",","projection_dim",")","`",",","where","`","N","`","is","the","total","number","of","examples","in","`","inputs_dataset","`",",","and","`","projection_dim","`","was","provided","in","initialization",".","Each","row","contains","the","vector","for","a","different","example","."]
code: def compute_intermediate_quantities(
      self,
      inputs_dataset: Union[Tuple[Any, ...], DataLoader],
      aggregate: bool = False,
      show_progress: bool = False,
      return_on_cpu: bool = True,
      test: bool = False,
  ) -> Tensor:
      r"""
      Computes "embedding" vectors for all examples in a single batch, or a
      `Dataloader` that yields batches. These embedding vectors are constructed so
      that the influence score of a training example on a test example is simply the
      dot-product of their corresponding vectors. In both cases, a batch should be
      small enough so that a backwards pass for a batch does not lead to
      out-of-memory errors.
  
      In more detail, the embedding vector for an example `x` is
      :math`\nabla_\theta L(x)' R`, where :math`R` is as defined in this class'
      description. Each element of `R` and :math`\nabla_\theta L(x)` lie in
      parameter-space. Therefore, if parameter-space were 1D, so that `R` were
      a 2D tensor whose columns are different elements in parameter-space, we would
      compute the embeddings for a batch by assembling :math`\nabla_\theta L(x)` for
      all examples `x` in the batch as rows in a 2D "batch parameter-gradient"
      tensor, and right-multiplying by `R`. However, parameter-space in this
      implementation is actually a tuple of tensors. So we do the analogous
      computation given this representation of parameter-space.
  
      If `aggregate` is True, the *sum* of the vectors for all examples is returned,
      instead of the vectors for each example. This can be useful for computing the
      influence of a given training example on the total loss over a validation
      dataset, because due to properties of the dot-product, this influence is the
      dot-product of the training example's vector with the sum of the vectors in the
      validation dataset. Also, by doing the sum aggregation within this method as
      opposed to outside of it (by computing all vectors for the validation dataset,
      then taking the sum) allows memory usage to be reduced.
  
      Args:
          inputs_dataset (Tuple, or DataLoader): Either a single tuple of any, or a
                  `DataLoader`, where each batch yielded is a tuple of any. In
                  either case, the tuple represents a single batch, where the last
                  element is assumed to be the labels for the batch. That is,
                  `model(*batch[0:-1])` produces the output for `model`, and
                  and `batch[-1]` are the labels, if any. Here, `model` is model
                  provided in initialization. This is the same assumption made for
                  each batch yielded by training dataset `train_dataset`.
          aggregate (bool): Whether to return the sum of the vectors for all
                  examples, as opposed to vectors for each example.
          show_progress (bool, optional): Computation of vectors can take a long
                  time if `inputs_dataset` represents many examples. If
                  `show_progress`is true, the progress of this computation will be
                  displayed. In particular, the number of batches for which
                  vectors have been computed will be displayed. It will try to
                  use tqdm if available for advanced features (e.g. time estimation).
                  Otherwise, it will fallback to a simple output of progress.
                  Default: False
          return_on_cpu (bool, optional): Whether to return the vectors on the cpu
                  (or if not, the gpu). If None, is set to the device that the model
                  is on.
                  Default: None
          test (bool, optional): Whether to compute the vectors using the loss
                  function `test_loss_fn` provided in initialization (instead of
                  `loss_fn`). This argument does not matter if `test_loss_fn` was
                  not provided, as in this case, `test_loss_fn` and `loss_fn` are the
                  same.
  
      Returns:
          intermediate_quantities (Tensor): This is a 2D tensor with shape
                  `(N, projection_dim)`, where `N` is the total number of examples in
                  `inputs_dataset`, and `projection_dim` was provided in
                  initialization. Each row contains the vector for a different
                  example.
      """
      # if `inputs_dataset` is not a `DataLoader`, turn it into one.
      inputs_dataset = _format_inputs_dataset(inputs_dataset)
  
      if show_progress:
          inputs_dataset = _progress_bar_constructor(
              self,
              inputs_dataset,
              "inputs_dataset",
              "intermediate quantities",
          )
  
      # infer model / data device through model. return device is same as that of
      # model unless explicitly specified
      if return_on_cpu is None:
          return_device = self.model_device
      else:
          return_device = (
              torch.device("cpu")
              if return_on_cpu
              else self.model_device
          )
  
      # choose the correct loss function and reduction type based on `test`
      loss_fn = self.test_loss_fn if test else self.loss_fn
      reduction_type = (
          self.test_reduction_type if test else self.reduction_type
      )
  
      # define a helper function that returns the embeddings for a batch
      def get_batch_embeddings(batch):
          # get gradient
          features, labels = tuple(batch[0:-1]), batch[-1]
          # `jacobians`` is a tensor of tuples. unlike parameters, however, the first
          # dimension is a batch dimension
          jacobians = _compute_jacobian_sample_wise_grads_per_batch(
              self, features, labels, loss_fn, reduction_type
          )
  
          # `jacobians`` contains the per-example parameters for a batch. this
          # function takes in `params`, a tuple of tensors representing a single
          # parameter setting, and for each example, computes the dot-product of its
          # per-example parameter with `params`. in other words, given `params`,
          # representing a basis vector, this function returns the coordinate of
          # each example in the batch along that basis. note that `jacobians` and
          # `params` are both tuple of tensors, with the same length. however, a
          # tensor in `jacobians` always has dimension 1 greater than the
          # corresponding tensor in `params`, because the tensors in `jacobians` have
          # a batch dimension (the 1st). to do this computation, the naive way would
          # be to convert `jacobians` to a list of tuple of tensors, and use
          # `_parameter_dot` to take the dot-product of each element in the list
          # with `params` to get a 1D tensor whose length is the batch size. however,
          # we can do the same computation without actually creating that list of
          # tuple of tensors by using broadcasting.
          def get_batch_coordinate(params):
              batch_coordinate = 0
              for _jacobians, param in zip(jacobians, params):
                  batch_coordinate += torch.sum(
                      _jacobians
                      * param.to(device=self.model_device).unsqueeze(0),
                      dim=tuple(range(1, len(_jacobians.shape))),
                  )
              return batch_coordinate.to(device=return_device)
  
          # to get the embedding for the batch, we get the coordinates for the batch
          # corresponding to one parameter in `R`. We do this for every parameter in
          # `R`, and then concatenate.
          return torch.stack(
              [get_batch_coordinate(params) for params in self.R],
              dim=1,
          )
  
      # using `get_batch_embeddings` and a helper, return all the vectors or their
      # sum, depending on `aggregate`
      return _get_dataset_embeddings_intermediate_quantities_influence_function(
          get_batch_embeddings,
          inputs_dataset,
          aggregate,
      )
  
imports: import functools
  from typing import Any, Callable, List, Optional, Tuple, Union
  import torch
  from captum._utils.gradient import _extract_parameters_from_layers
  from captum.influence._utils.common import _compute_batch_loss_influence_function_base, _compute_jacobian_sample_wise_grads_per_batch, _dataset_fn, _format_inputs_dataset, _functional_call, _get_k_most_influential_helper, _influence_batch_intermediate_quantities_influence_function, _influence_helper_intermediate_quantities_influence_function, _influence_route_to_helpers, _load_flexible_state_dict, _parameter_add, _parameter_dot, _parameter_linear_combination, _parameter_multiply, _parameter_to, _params_to_names, _progress_bar_constructor, _self_influence_helper_intermediate_quantities_influence_function, _top_eigen, KMostInfluentialResults
  from captum.log import log_usage
  from torch import Tensor
  from torch.nn import Module
  from torch.utils.data import DataLoader, Dataset
  from tqdm import tqdm
  from .influence_function import _get_dataset_embeddings_intermediate_quantities_influence_function, InfluenceFunctionBase, IntermediateQuantitiesInfluenceFunction
label_desc: Use image node_id 3 for calling the ArnoldiInfluenceFunction obj's underlying member method code with example usage: obj.compute_intermediate_quantities(inputs_dataset, aggregate, show_progress, return_on_cpu, test) and returns: _get_dataset_embeddings_intermediate_quantities_influence_function, torch, batch_coordinate