docstr: This is a helper with 2 modes. For both modes, it first generates a matrix
  with `A` with specified eigenvalues.
  
  When mode is 'arnoldi', it checks that `_parameter_arnoldi` is correct.
  In particular, it checks that the top-`k` eigenvalues of the restriction
  of `A` to a Krylov subspace (the `H` returned by `_parameter_arnoldi`)
  agree with those of the original matrix. This is a property we expect of the
  Arnoldi iteration that `_parameter_arnoldi` implements.
  
  When mode is 'distill', it checks that `_parameter_distill` is correct. In
  particular, it checks that the eigenvectors corresponding to the top
  eigenvalues it returns agree with the top eigenvectors of `A`. This is the
  property we require of `distill`, because we use the top eigenvectors (and
  eigenvalues) of (implicitly-defined) `A` to calculate a low-rank approximation
  of its inverse.
docstr_tok: ["This","is","a","helper","with","2","modes",".","For","both","modes",",","it","first","generates","a","matrix","with","`","A","`","with","specified","eigenvalues",".","When","mode","is","'arnoldi","'",",","it","checks","that","`","_parameter_arnoldi","`","is","correct",".","In","particular",",","it","checks","that","the","top-","`","k","`","eigenvalues","of","the","restriction","of","`","A","`","to","a","Krylov","subspace","(","the","`","H","`","returned","by","`","_parameter_arnoldi","`",")","agree","with","those","of","the","original","matrix",".","This","is","a","property","we","expect","of","the","Arnoldi","iteration","that","`","_parameter_arnoldi","`","implements",".","When","mode","is","'distill","'",",","it","checks","that","`","_parameter_distill","`","is","correct",".","In","particular",",","it","checks","that","the","eigenvectors","corresponding","to","the","top","eigenvalues","it","returns","agree","with","the","top","eigenvectors","of","`","A","`",".","This","is","the","property","we","require","of","`","distill","`",",","because","we","use","the","top","eigenvectors","(","and","eigenvalues",")","of","(","implicitly-defined",")","`","A","`","to","calculate","a","low-rank","approximation","of","its","inverse","."]
code: def _test_parameter_arnoldi_and_distill(
      self,
      mode: str,
      symmetric: bool,
      eigenvalues: Tensor,
      k: int,
      arnoldi_dim: int,
      param_shape: List[Tuple],
  ):
      """
      This is a helper with 2 modes. For both modes, it first generates a matrix
      with `A` with specified eigenvalues.
  
      When mode is 'arnoldi', it checks that `_parameter_arnoldi` is correct.
      In particular, it checks that the top-`k` eigenvalues of the restriction
      of `A` to a Krylov subspace (the `H` returned by `_parameter_arnoldi`)
      agree with those of the original matrix. This is a property we expect of the
      Arnoldi iteration that `_parameter_arnoldi` implements.
  
      When mode is 'distill', it checks that `_parameter_distill` is correct. In
      particular, it checks that the eigenvectors corresponding to the top
      eigenvalues it returns agree with the top eigenvectors of `A`. This is the
      property we require of `distill`, because we use the top eigenvectors (and
      eigenvalues) of (implicitly-defined) `A` to calculate a low-rank approximation
      of its inverse.
      """
      # generate matrix `A` with specified eigenvalues
      A = (
          generate_symmetric_matrix_given_eigenvalues(eigenvalues)
          if symmetric
          else generate_assymetric_matrix_given_eigenvalues(eigenvalues)
      )
  
      # create the matrix-vector multiplication function that `_parameter_arnoldi`
      # expects that represents multiplication by `A`.
      # since the vector actually needs to be a tuple of tensors, we
      # specify the dimensions of that tuple of tensors. the function then
      # flattens the vector, multiplies it by the generated matrix, and then
      # unflattens the result
      _unflatten_params = _unflatten_params_factory(param_shape)
  
      def _param_matmul(params: Tuple[Tensor]):
          return _unflatten_params(
              torch.matmul(A, _flatten_params(params))
          )
  
      # generate `b` and call `_parameter_arnoldi`
      b = tuple(torch.randn(shape) for shape in param_shape)
      qs, H = _parameter_arnoldi(
          _param_matmul,
          b,
          arnoldi_dim,
          1e-3,
          torch.device("cpu"),
          False,
      )
  
      assertTensorAlmostEqual(
          self,
          _flatten_params(_unflatten_params(_flatten_params(b))),
          _flatten_params(b),
          1e-5,
          "max",
      )
  
      # compute the eigenvalues / eigenvectors of `A` and `H`. we use `eig` since
      # each matrix may not be symmetric. since `eig` does not sort by eigenvalues,
      # need to manually do it. also get rid of last column of H, since
      # it is not part of the decomposition
      vs_A, ls_A = _eig_helper(A)
      vs_H, ls_H = _eig_helper(H[:-1])
  
      if mode == "arnoldi":
          # compare the top-`k` eigenvalue of the two matrices
          assertTensorAlmostEqual(
              self, vs_H[-k:], vs_A[-k:], 1e-3, "max"
          )
      elif mode == "distill":
          # use `distill` to compute top-`k` eigenvectors of `H` in the original
          # basis. then check if they are actually eigenvectors
          vs_H_standard, ls_H_standard = _parameter_distill(
              qs, H, k, 0, 0
          )
  
          for l_H_standard, v_A in zip(ls_H_standard[-k:], vs_A[-k:]):
              l_H_standard_flattened = _flatten_params(
                  l_H_standard
              )  # .real
              expected = v_A * l_H_standard_flattened
              actual = torch.matmul(A, l_H_standard_flattened)
              # tol copied from original code
              assert torch.norm(expected - actual) < 1e-2
  
          # check that the top-`k` eigenvalues of `A` as computed by
          # `_parameters_distill` are similar to those computed on `A` directly
          for v_H_standard, v_A in zip(vs_H_standard[-k:], vs_A[-k:]):
              # tol copied from original code
              assert abs(v_H_standard - v_A) < 5
  
          if False:
              # code from original paper does not do this test, so skip for now
              # use `distill`` to get top-`k` eigenvectors of `H` in the original
              # basis, and compare with the top-`k` eigenvectors of `A`. need to
              # flatten those from `distill` to compare
              _, ls_H_standard = _parameter_distill(qs, H, k, 0, 0)
              for l_H_standard, l_A in zip(ls_H_standard, ls_A):
                  # print(l_A)
                  # print(flatten_unflattener.flatten(l_H_standard).real)
                  l_H_standard_flattened /= torch.norm(
                      l_H_standard_flattened
                  )
                  assertTensorAlmostEqual(
                      self,
                      _flatten_params(l_H_standard).real,
                      l_A.real,
                      1e-2,
                      "max",
                  )
  
imports: import tempfile
  from typing import Callable, List, Tuple, Union
  import torch
  import torch.nn
  from captum.influence._core.arnoldi_influence_function import _parameter_arnoldi, _parameter_distill, ArnoldiInfluenceFunction
  from captum.influence._core.influence_function import NaiveInfluenceFunction
  from captum.influence._utils.common import _eig_helper, _flatten_params, _top_eigen, _unflatten_params_factory
  from parameterized import parameterized
  from tests.helpers.basic import assertTensorAlmostEqual, BaseTest
  from tests.influence._utils.common import _format_batch_into_tuple, build_test_name_func, DataInfluenceConstructor, ExplicitDataset, generate_assymetric_matrix_given_eigenvalues, generate_symmetric_matrix_given_eigenvalues, get_random_model_and_data, UnpackDataset, USE_GPU_LIST
  from torch import Tensor
  from torch.utils.data import DataLoader
label_desc: Use image node_id 4 for calling the TestArnoldiInfluence obj's underlying member method code with example usage: obj._test_parameter_arnoldi_and_distill(mode, symmetric, eigenvalues, k, arnoldi_dim, param_shape) and returns: _unflatten_params