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: Returns the `R` described in the documentation for
  `ArnoldiInfluenceFunction`. The returned `R` represents a set of
  parameters in parameter space. However, since this implementation does *not*
  flatten parameters, each of those parameters is represented as a tuple of
  tensors. Therefore, `R` is represented as a list of tuple of tensors, and
  can be viewed as a linear function that takes in a tuple of tensors
  (representing a parameter), and returns a vector, where the i-th entry is
  the dot-product (as it would be defined over tuple of tensors) of the parameter
  (i.e. the input to the linear function) with the i-th entry of `R`.
  
  Can specify that projection should always be saved on cpu. if so, gradients are
  always moved to same device as projections before multiplying (moving
  projections to gpu when multiplying would defeat the purpose of moving them to
  cpu to save gpu memory).
  
  Returns:
      R (list of tuple of tensors): List of tuple of tensors of length
              `projection_dim` (initialization argument). Each element
              corresponds to a parameter in parameter-space, is represented as a
              tuple of tensors, and together, define a projection that can be
              applied to parameters (represented as tuple of tensors).
docstr_tok: ["Returns","the","`","R","`","described","in","the","documentation","for","`","ArnoldiInfluenceFunction","`",".","The","returned","`","R","`","represents","a","set","of","parameters","in","parameter","space",".","However",",","since","this","implementation","does","*","not","*","flatten","parameters",",","each","of","those","parameters","is","represented","as","a","tuple","of","tensors",".","Therefore",",","`","R","`","is","represented","as","a","list","of","tuple","of","tensors",",","and","can","be","viewed","as","a","linear","function","that","takes","in","a","tuple","of","tensors","(","representing","a","parameter",")",",","and","returns","a","vector",",","where","the","i-th","entry","is","the","dot-product","(","as","it","would","be","defined","over","tuple","of","tensors",")","of","the","parameter","(","i.e",".","the","input","to","the","linear","function",")","with","the","i-th","entry","of","`","R","`",".","Can","specify","that","projection","should","always","be","saved","on","cpu",".","if","so",",","gradients","are","always","moved","to","same","device","as","projections","before","multiplying","(","moving","projections","to","gpu","when","multiplying","would","defeat","the","purpose","of","moving","them","to","cpu","to","save","gpu","memory",")",".","Returns",":","R","(","list","of","tuple","of","tensors",")",":","List","of","tuple","of","tensors","of","length","`","projection_dim","`","(","initialization","argument",")",".","Each","element","corresponds","to","a","parameter","in","parameter-space",",","is","represented","as","a","tuple","of","tensors",",","and","together",",","define","a","projection","that","can","be","applied","to","parameters","(","represented","as","tuple","of","tensors",")","."]
code: def _retrieve_projections_arnoldi_influence_function(
      self,
      dataloader: DataLoader,
      projection_on_cpu: bool,
      show_progress: bool,
  ) -> List[Tuple[Tensor, ...]]:
      """
  
      Returns the `R` described in the documentation for
      `ArnoldiInfluenceFunction`. The returned `R` represents a set of
      parameters in parameter space. However, since this implementation does *not*
      flatten parameters, each of those parameters is represented as a tuple of
      tensors. Therefore, `R` is represented as a list of tuple of tensors, and
      can be viewed as a linear function that takes in a tuple of tensors
      (representing a parameter), and returns a vector, where the i-th entry is
      the dot-product (as it would be defined over tuple of tensors) of the parameter
      (i.e. the input to the linear function) with the i-th entry of `R`.
  
      Can specify that projection should always be saved on cpu. if so, gradients are
      always moved to same device as projections before multiplying (moving
      projections to gpu when multiplying would defeat the purpose of moving them to
      cpu to save gpu memory).
  
      Returns:
          R (list of tuple of tensors): List of tuple of tensors of length
                  `projection_dim` (initialization argument). Each element
                  corresponds to a parameter in parameter-space, is represented as a
                  tuple of tensors, and together, define a projection that can be
                  applied to parameters (represented as tuple of tensors).
      """
      # create function that computes hessian-vector product, given a vector
      # represented as a tuple of tensors
  
      # first figure out names of params that require gradients. this is need to
      # create that function, as it replaces params based on their names
      params = tuple(
          self.model.parameters()
          if self.layer_modules is None
          else _extract_parameters_from_layers(self.layer_modules)
      )
      # the same position in `params` and `param_names` correspond to each other
      param_names = _params_to_names(params, self.model)
  
      # get factory that given a batch, returns a function that given params as
      # tuple of tensors, returns loss over the batch
      def tensor_tuple_loss_given_batch(batch):
          def tensor_tuple_loss(*params):
              # `params` is a tuple of tensors, and assumed to be order specified by
              # `param_names`
              features, labels = tuple(batch[0:-1]), batch[-1]
  
              _output = _functional_call(
                  self.model, dict(zip(param_names, params)), features
              )
  
              # compute the total loss for the batch, adjusting the output of
              # `self.loss_fn` based on `self.reduction_type`
              return _compute_batch_loss_influence_function_base(
                  self.loss_fn, _output, labels, self.reduction_type
              )
  
          return tensor_tuple_loss
  
      # define function that given batch and vector, returns HVP of loss using the
      # batch and vector
      def batch_HVP(batch, v):
          tensor_tuple_loss = tensor_tuple_loss_given_batch(batch)
          return torch.autograd.functional.hvp(
              tensor_tuple_loss, params, v=v
          )[1]
  
      # define function that returns HVP of loss over `dataloader`, given a
      # specified vector
      def HVP(v):
          _hvp = None
  
          _dataloader = dataloader
          if show_progress:
              _dataloader = tqdm(
                  dataloader, desc="processing `hessian_dataset` batch"
              )
  
          # the HVP of loss using the entire `dataloader` is the sum of the
          # per-batch HVP's
          return _dataset_fn(_dataloader, batch_HVP, _parameter_add, v)
  
          for batch in _dataloader:
              hvp = batch_HVP(batch, v)
              if _hvp is None:
                  _hvp = hvp
              else:
                  _hvp = _parameter_add(_hvp, hvp)
          return _hvp
  
      # now that can compute the hessian-vector product (of loss over `dataloader`),
      # can perform arnoldi iteration
  
      # we always perform the HVP computations on the device where the model is.
      # effectively this means we do the computations on gpu if available. this
      # is necessary because the HVP is computationally expensive.
  
      # get initial random vector, and place it on the same device as the model.
      # `_parameter_arnoldi` needs to know which device the model is on, and
      # will infer it through the device of this random vector
      b = _parameter_to(
          tuple(torch.randn_like(param) for param in params),
          device=self.model_device,
      )
  
      # perform the arnoldi iteration, see its documentation for what its return
      # values are.  note that `H` is *not* the Hessian.
      qs, H = _parameter_arnoldi(
          HVP,
          b,
          self.arnoldi_dim,
          self.arnoldi_tol,
          torch.device("cpu")
          if projection_on_cpu
          else self.model_device,
          show_progress,
      )
  
      # `ls`` and `vs`` are (approximately) the top eigenvalues / eigenvectors of the
      # matrix used (implicitly) to compute Hessian-vector products by the `HVP`
      # input to `_parameter_arnoldi`. this matrix is the Hessian of the loss,
      # summed over the examples in `dataloader`. note that because the vectors in
      # the Hessian-vector product are actually tuples of tensors representing
      # parameters, `vs`` is a list of tuples of tensors.  note that here, `H` is
      # *not* the Hessian (`qs` and `H` together define the Krylov subspace of the
      # Hessian)
  
      ls, vs = _parameter_distill(
          qs,
          H,
          self.projection_dim,
          self.hessian_reg,
          self.hessian_inverse_tol,
      )
  
      # if `vs` were a 2D tensor whose columns contain the top eigenvectors of the
      # aforementioned hessian, then `R` would be `vs @ torch.diag(ls ** -0.5)`, i.e.
      # scaling each column of `vs` by the corresponding entry in `ls ** -0.5`.
      # however, since `vs` is instead a list of tuple of tensors, `R` should be
      # a list of tuple of tensors, where each entry in the list is scaled by the
      # corresponding entry in `ls ** 0.5`, which we first compute.
      ls = (1.0 / ls) ** 0.5
  
      # then, scale each entry in `vs` by the corresponding entry in `ls ** 0.5`
      # since each entry in `vs` is a tuple of tensors, we use a helper function
      # that takes in a tuple of tensors, and a scalar, and multiplies every tensor
      # by the scalar.
      return [_parameter_multiply(v, l) for (v, l) in zip(vs, ls)]
  
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 2 for calling the ArnoldiInfluenceFunction obj's underlying member method code with example usage: obj._retrieve_projections_arnoldi_influence_function(dataloader, projection_on_cpu, show_progress) and returns: unknown, tensor_tuple_loss, torch, _dataset_fn, _hvp, _compute_batch_loss_influence_function_base