class_docstr: `InfluenceFunctionBase` is a base class for implementations which compute the
  influence score as defined in the paper "Understanding Black-box Predictions via
  Influence Functions" (https://arxiv.org/pdf/1703.04730.pdf). This "infinitesimal"
  influence score approximately answers the question if a given training example
  were infinitesimally down-weighted and the model re-trained to optimality, how much
  would the loss on a given test example change. Mathematically, the aforementioned
  influence score is given by :math`\nabla_\theta L(x)' H^{-1} \nabla_\theta L(z)`,
  where :math`\nabla_\theta L(x)` is the gradient of the loss, considering only
  training example :math`x` with respect to (a subset of) model parameters
  :math`\theta`, :math`\nabla_\theta L(z)` is the analogous quantity for a test
  example :math`z`, and :math`H` is the Hessian of the (subset of) model parameters
  at a given model checkpoint. "Subset of model parameters" refers to the parameters
  specified by the `layers` initialization argument; for computational purposes,
  we may only consider the gradients / Hessian involving parameters in a subset of
  the model's layers. This is a commonly-taken approach in the research literature.
  
  There can be multiple implementations of this class, because although the paper
  defines a particular "infinitesimal" kind of influence score, there can be multiple
  ways to compute it, each with different levels of accuracy / scalability.
docstr: Args:
      model (torch.nn.Module): An instance of pytorch model. This model should
              define all of its layers as attributes of the model.
      train_dataset (torch.utils.data.Dataset or torch.utils.data.DataLoader):
              In the `influence` method, we either compute the influence score of
              training examples on examples in a test batch, or self influence
              scores for those training examples, depending on which mode is used.
              This argument represents the training dataset containing those
              training examples. In order to compute those influence scores, we
              will create a Pytorch DataLoader yielding batches of training
              examples that is then used for processing. If this argument is
              already a Pytorch Dataloader, that DataLoader can be directly
              used for processing. If it is instead a Pytorch Dataset, we will
              create a DataLoader using it, with batch size specified by
              `batch_size`. For efficiency purposes, the batch size of the
              DataLoader used for processing should be as large as possible, but
              not too large, so that certain intermediate quantities created
              from a batch still fit in memory. Therefore, if
              `train_dataset` is a Dataset, `batch_size` should be large.
              If `train_dataset` was already a DataLoader to begin with,
              it should have been constructed to have a large batch size. It is
              assumed that the Dataloader (regardless of whether it is created
              from a Pytorch Dataset or not) yields tuples. For a `batch` that is
              yielded, of length `L`, it is assumed that the forward function of
              `model` accepts `L-1` arguments, and the last element of `batch` is
              the label. In other words, `model(*batch[:-1])` gives the output of
              `model`, and `batch[-1]` are the labels for the batch.
      checkpoint (str): The path to the checkpoint used to compute influence
              scores.
      checkpoints_load_func (Callable, optional): The function to load a saved
              checkpoint into a model to update its parameters, and get the
              learning rate if it is saved. By default uses a utility to load a
              model saved as a state dict.
              Default: _load_flexible_state_dict
      layers (list[str] or None, optional): A list of layer names for which
              gradients should be computed. If `layers` is None, gradients will
              be computed for all layers. Otherwise, they will only be computed
              for the layers specified in `layers`.
              Default: None
      loss_fn (Callable, optional): The loss function applied to model. There
              are two options for the return type of `loss_fn`. First, `loss_fn`
              can be a "per-example" loss function - returns a 1D Tensor of
              losses for each example in a batch. `nn.BCELoss(reduction="none")`
              would be an "per-example" loss function. Second, `loss_fn` can be
              a "reduction" loss function that reduces the per-example losses,
              in a batch, and returns a single scalar Tensor. For this option,
              the reduction must be the *sum* or the *mean* of the per-example
              losses. For instance, `nn.BCELoss(reduction="sum")` is acceptable.
              Note for the first option, the `sample_wise_grads_per_batch`
              argument must be False, and for the second option,
              `sample_wise_grads_per_batch` must be True.  Also note that for
              the second option, if `loss_fn` has no "reduction" attribute,
              the implementation assumes that the reduction is the *sum* of the
              per-example losses.  If this is not the case, i.e. the reduction
              is the *mean*, please set the "reduction" attribute of `loss_fn`
              to "mean", i.e. `loss_fn.reduction = "mean"`.
      batch_size (int or None, optional): Batch size of the DataLoader created to
              iterate through `train_dataset` and `hessian_dataset`, if they are
              of type `Dataset`. `batch_size` should be chosen as large as
              possible so that a backwards pass on a batch still fits in memory.
              If `train_dataset` and `hessian_dataset`are both of type
              `DataLoader`, then `batch_size` is ignored as an argument.
              Default: 1
      hessian_dataset (Dataset or Dataloader, optional): The influence score and
              self-influence scores this implementation calculates are defined in
              terms of the Hessian, i.e. the second-derivative of the model
              parameters. This argument provides the dataset used for calculating
              the Hessian. It should be smaller than `train_dataset`, which
              is the dataset whose examples we want the influence of. If not
              provided or none, it will be assumed to be the same as
              `train_dataset`.
              Default: None
      test_loss_fn (Callable, optional): In some cases, one may want to use a
              separate loss functions for training examples, i.e. those in
              `train_dataset`, and for test examples, i.e. those
              represented by the `inputs` and `targets` arguments to the
              `influence` method. For example, if one wants to calculate the
              influence score of a training example on a test example's
              prediction for a fixed class, `test_loss_fn` could map from the
              logits for all classes to the logits for a fixed class.
              `test_loss_fn` needs satisfy the same constraints as `loss_fn`.
              Thus, the same checks that we apply to `loss_fn` are also applied
              to `test_loss_fn`, if the latter is provided. Note that the
              constraints on both `loss_fn` and `test_loss_fn` both depend on
              `sample_wise_grads_per_batch`. This means `loss_fn` and
              `test_loss_fn` must either both be "per-example"  loss functions,
              or both be "reduction" loss functions. If not provided, the loss
              function for test examples is assumed to be the same as the loss
              function for training examples, i.e. `loss_fn`.
              Default: None
      sample_wise_grads_per_batch (bool, optional): PyTorch's native gradient
              computations w.r.t. model parameters aggregates the results for a
              batch and does not allow to access sample-wise gradients w.r.t.
              model parameters. This forces us to iterate over each sample in
              the batch if we want sample-wise gradients which is computationally
              inefficient. We offer an implementation of batch-wise gradient
              computations w.r.t. to model parameters which is computationally
              more efficient. This implementation can be enabled by setting the
              `sample_wise_grad_per_batch` argument to `True`, and should be
              enabled if and only if the `loss_fn` argument is a "reduction" loss
              function. For example, `nn.BCELoss(reduction="sum")` would be a
              valid `loss_fn` if this implementation is enabled (see
              documentation for `loss_fn` for more details). Note that our
              current implementation enables batch-wise gradient computations
              only for a limited number of PyTorch nn.Modules: Conv2D and Linear.
              This list will be expanded in the near future.  Therefore, please
              do not enable this implementation if gradients will be computed
              for other kinds of layers.
              Default: False
docstr_tok: ["Args",":","model","(","torch.nn.Module",")",":","An","instance","of","pytorch","model",".","This","model","should","define","all","of","its","layers","as","attributes","of","the","model",".","train_dataset","(","torch.utils.data.Dataset","or","torch.utils.data.DataLoader",")",":","In","the","`","influence","`","method",",","we","either","compute","the","influence","score","of","training","examples","on","examples","in","a","test","batch",",","or","self","influence","scores","for","those","training","examples",",","depending","on","which","mode","is","used",".","This","argument","represents","the","training","dataset","containing","those","training","examples",".","In","order","to","compute","those","influence","scores",",","we","will","create","a","Pytorch","DataLoader","yielding","batches","of","training","examples","that","is","then","used","for","processing",".","If","this","argument","is","already","a","Pytorch","Dataloader",",","that","DataLoader","can","be","directly","used","for","processing",".","If","it","is","instead","a","Pytorch","Dataset",",","we","will","create","a","DataLoader","using","it",",","with","batch","size","specified","by","`","batch_size","`",".","For","efficiency","purposes",",","the","batch","size","of","the","DataLoader","used","for","processing","should","be","as","large","as","possible",",","but","not","too","large",",","so","that","certain","intermediate","quantities","created","from","a","batch","still","fit","in","memory",".","Therefore",",","if","`","train_dataset","`","is","a","Dataset",",","`","batch_size","`","should","be","large",".","If","`","train_dataset","`","was","already","a","DataLoader","to","begin","with",",","it","should","have","been","constructed","to","have","a","large","batch","size",".","It","is","assumed","that","the","Dataloader","(","regardless","of","whether","it","is","created","from","a","Pytorch","Dataset","or","not",")","yields","tuples",".","For","a","`","batch","`","that","is","yielded",",","of","length","`","L","`",",","it","is","assumed","that","the","forward","function","of","`","model","`","accepts","`","L-1","`","arguments",",","and","the","last","element","of","`","batch","`","is","the","label",".","In","other","words",",","`","model","(","*","batch","[",":","-1","]",")","`","gives","the","output","of","`","model","`",",","and","`","batch","[","-1","]","`","are","the","labels","for","the","batch",".","checkpoint","(","str",")",":","The","path","to","the","checkpoint","used","to","compute","influence","scores",".","checkpoints_load_func","(","Callable",",","optional",")",":","The","function","to","load","a","saved","checkpoint","into","a","model","to","update","its","parameters",",","and","get","the","learning","rate","if","it","is","saved",".","By","default","uses","a","utility","to","load","a","model","saved","as","a","state","dict",".","Default",":","_load_flexible_state_dict","layers","(","list","[","str","]","or","None",",","optional",")",":","A","list","of","layer","names","for","which","gradients","should","be","computed",".","If","`","layers","`","is","None",",","gradients","will","be","computed","for","all","layers",".","Otherwise",",","they","will","only","be","computed","for","the","layers","specified","in","`","layers","`",".","Default",":","None","loss_fn","(","Callable",",","optional",")",":","The","loss","function","applied","to","model",".","There","are","two","options","for","the","return","type","of","`","loss_fn","`",".","First",",","`","loss_fn","`","can","be","a","``","per-example","''","loss","function","-","returns","a","1D","Tensor","of","losses","for","each","example","in","a","batch",".","`","nn.BCELoss","(","reduction=","''","none","''",")","`","would","be","an","``","per-example","''","loss","function",".","Second",",","`","loss_fn","`","can","be","a","``","reduction","''","loss","function","that","reduces","the","per-example","losses",",","in","a","batch",",","and","returns","a","single","scalar","Tensor",".","For","this","option",",","the","reduction","must","be","the","*","sum","*","or","the","*","mean","*","of","the","per-example","losses",".","For","instance",",","`","nn.BCELoss","(","reduction=","''","sum","''",")","`","is","acceptable",".","Note","for","the","first","option",",","the","`","sample_wise_grads_per_batch","`","argument","must","be","False",",","and","for","the","second","option",",","`","sample_wise_grads_per_batch","`","must","be","True",".","Also","note","that","for","the","second","option",",","if","`","loss_fn","`","has","no","``","reduction","''","attribute",",","the","implementation","assumes","that","the","reduction","is","the","*","sum","*","of","the","per-example","losses",".","If","this","is","not","the","case",",","i.e",".","the","reduction","is","the","*","mean","*",",","please","set","the","``","reduction","''","attribute","of","`","loss_fn","`","to","``","mean","''",",","i.e",".","`","loss_fn.reduction","=","``","mean","''","`",".","batch_size","(","int","or","None",",","optional",")",":","Batch","size","of","the","DataLoader","created","to","iterate","through","`","train_dataset","`","and","`","hessian_dataset","`",",","if","they","are","of","type","`","Dataset","`",".","`","batch_size","`","should","be","chosen","as","large","as","possible","so","that","a","backwards","pass","on","a","batch","still","fits","in","memory",".","If","`","train_dataset","`","and","`","hessian_dataset","`","are","both","of","type","`","DataLoader","`",",","then","`","batch_size","`","is","ignored","as","an","argument",".","Default",":","1","hessian_dataset","(","Dataset","or","Dataloader",",","optional",")",":","The","influence","score","and","self-influence","scores","this","implementation","calculates","are","defined","in","terms","of","the","Hessian",",","i.e",".","the","second-derivative","of","the","model","parameters",".","This","argument","provides","the","dataset","used","for","calculating","the","Hessian",".","It","should","be","smaller","than","`","train_dataset","`",",","which","is","the","dataset","whose","examples","we","want","the","influence","of",".","If","not","provided","or","none",",","it","will","be","assumed","to","be","the","same","as","`","train_dataset","`",".","Default",":","None","test_loss_fn","(","Callable",",","optional",")",":","In","some","cases",",","one","may","want","to","use","a","separate","loss","functions","for","training","examples",",","i.e",".","those","in","`","train_dataset","`",",","and","for","test","examples",",","i.e",".","those","represented","by","the","`","inputs","`","and","`","targets","`","arguments","to","the","`","influence","`","method",".","For","example",",","if","one","wants","to","calculate","the","influence","score","of","a","training","example","on","a","test","example's","prediction","for","a","fixed","class",",","`","test_loss_fn","`","could","map","from","the","logits","for","all","classes","to","the","logits","for","a","fixed","class",".","`","test_loss_fn","`","needs","satisfy","the","same","constraints","as","`","loss_fn","`",".","Thus",",","the","same","checks","that","we","apply","to","`","loss_fn","`","are","also","applied","to","`","test_loss_fn","`",",","if","the","latter","is","provided",".","Note","that","the","constraints","on","both","`","loss_fn","`","and","`","test_loss_fn","`","both","depend","on","`","sample_wise_grads_per_batch","`",".","This","means","`","loss_fn","`","and","`","test_loss_fn","`","must","either","both","be","``","per-example","''","loss","functions",",","or","both","be","``","reduction","''","loss","functions",".","If","not","provided",",","the","loss","function","for","test","examples","is","assumed","to","be","the","same","as","the","loss","function","for","training","examples",",","i.e",".","`","loss_fn","`",".","Default",":","None","sample_wise_grads_per_batch","(","bool",",","optional",")",":","PyTorch","'s","native","gradient","computations","w.r.t",".","model","parameters","aggregates","the","results","for","a","batch","and","does","not","allow","to","access","sample-wise","gradients","w.r.t",".","model","parameters",".","This","forces","us","to","iterate","over","each","sample","in","the","batch","if","we","want","sample-wise","gradients","which","is","computationally","inefficient",".","We","offer","an","implementation","of","batch-wise","gradient","computations","w.r.t",".","to","model","parameters","which","is","computationally","more","efficient",".","This","implementation","can","be","enabled","by","setting","the","`","sample_wise_grad_per_batch","`","argument","to","`","True","`",",","and","should","be","enabled","if","and","only","if","the","`","loss_fn","`","argument","is","a","``","reduction","''","loss","function",".","For","example",",","`","nn.BCELoss","(","reduction=","''","sum","''",")","`","would","be","a","valid","`","loss_fn","`","if","this","implementation","is","enabled","(","see","documentation","for","`","loss_fn","`","for","more","details",")",".","Note","that","our","current","implementation","enables","batch-wise","gradient","computations","only","for","a","limited","number","of","PyTorch","nn.Modules",":","Conv2D","and","Linear",".","This","list","will","be","expanded","in","the","near","future",".","Therefore",",","please","do","not","enable","this","implementation","if","gradients","will","be","computed","for","other","kinds","of","layers",".","Default",":","False"]
code: def __init__(
      self,
      model: Module,
      train_dataset: Union[Dataset, DataLoader],
      checkpoint: str,
      checkpoints_load_func: Callable = _load_flexible_state_dict,
      layers: Optional[List[str]] = None,
      loss_fn: Optional[Union[Module, Callable]] = None,
      batch_size: Union[int, None] = 1,
      hessian_dataset: Optional[Union[Dataset, DataLoader]] = None,
      test_loss_fn: Optional[Union[Module, Callable]] = None,
      sample_wise_grads_per_batch: bool = False,
  ) -> None:
      """
      Args:
          model (torch.nn.Module): An instance of pytorch model. This model should
                  define all of its layers as attributes of the model.
          train_dataset (torch.utils.data.Dataset or torch.utils.data.DataLoader):
                  In the `influence` method, we either compute the influence score of
                  training examples on examples in a test batch, or self influence
                  scores for those training examples, depending on which mode is used.
                  This argument represents the training dataset containing those
                  training examples. In order to compute those influence scores, we
                  will create a Pytorch DataLoader yielding batches of training
                  examples that is then used for processing. If this argument is
                  already a Pytorch Dataloader, that DataLoader can be directly
                  used for processing. If it is instead a Pytorch Dataset, we will
                  create a DataLoader using it, with batch size specified by
                  `batch_size`. For efficiency purposes, the batch size of the
                  DataLoader used for processing should be as large as possible, but
                  not too large, so that certain intermediate quantities created
                  from a batch still fit in memory. Therefore, if
                  `train_dataset` is a Dataset, `batch_size` should be large.
                  If `train_dataset` was already a DataLoader to begin with,
                  it should have been constructed to have a large batch size. It is
                  assumed that the Dataloader (regardless of whether it is created
                  from a Pytorch Dataset or not) yields tuples. For a `batch` that is
                  yielded, of length `L`, it is assumed that the forward function of
                  `model` accepts `L-1` arguments, and the last element of `batch` is
                  the label. In other words, `model(*batch[:-1])` gives the output of
                  `model`, and `batch[-1]` are the labels for the batch.
          checkpoint (str): The path to the checkpoint used to compute influence
                  scores.
          checkpoints_load_func (Callable, optional): The function to load a saved
                  checkpoint into a model to update its parameters, and get the
                  learning rate if it is saved. By default uses a utility to load a
                  model saved as a state dict.
                  Default: _load_flexible_state_dict
          layers (list[str] or None, optional): A list of layer names for which
                  gradients should be computed. If `layers` is None, gradients will
                  be computed for all layers. Otherwise, they will only be computed
                  for the layers specified in `layers`.
                  Default: None
          loss_fn (Callable, optional): The loss function applied to model. There
                  are two options for the return type of `loss_fn`. First, `loss_fn`
                  can be a "per-example" loss function - returns a 1D Tensor of
                  losses for each example in a batch. `nn.BCELoss(reduction="none")`
                  would be an "per-example" loss function. Second, `loss_fn` can be
                  a "reduction" loss function that reduces the per-example losses,
                  in a batch, and returns a single scalar Tensor. For this option,
                  the reduction must be the *sum* or the *mean* of the per-example
                  losses. For instance, `nn.BCELoss(reduction="sum")` is acceptable.
                  Note for the first option, the `sample_wise_grads_per_batch`
                  argument must be False, and for the second option,
                  `sample_wise_grads_per_batch` must be True.  Also note that for
                  the second option, if `loss_fn` has no "reduction" attribute,
                  the implementation assumes that the reduction is the *sum* of the
                  per-example losses.  If this is not the case, i.e. the reduction
                  is the *mean*, please set the "reduction" attribute of `loss_fn`
                  to "mean", i.e. `loss_fn.reduction = "mean"`.
          batch_size (int or None, optional): Batch size of the DataLoader created to
                  iterate through `train_dataset` and `hessian_dataset`, if they are
                  of type `Dataset`. `batch_size` should be chosen as large as
                  possible so that a backwards pass on a batch still fits in memory.
                  If `train_dataset` and `hessian_dataset`are both of type
                  `DataLoader`, then `batch_size` is ignored as an argument.
                  Default: 1
          hessian_dataset (Dataset or Dataloader, optional): The influence score and
                  self-influence scores this implementation calculates are defined in
                  terms of the Hessian, i.e. the second-derivative of the model
                  parameters. This argument provides the dataset used for calculating
                  the Hessian. It should be smaller than `train_dataset`, which
                  is the dataset whose examples we want the influence of. If not
                  provided or none, it will be assumed to be the same as
                  `train_dataset`.
                  Default: None
          test_loss_fn (Callable, optional): In some cases, one may want to use a
                  separate loss functions for training examples, i.e. those in
                  `train_dataset`, and for test examples, i.e. those
                  represented by the `inputs` and `targets` arguments to the
                  `influence` method. For example, if one wants to calculate the
                  influence score of a training example on a test example's
                  prediction for a fixed class, `test_loss_fn` could map from the
                  logits for all classes to the logits for a fixed class.
                  `test_loss_fn` needs satisfy the same constraints as `loss_fn`.
                  Thus, the same checks that we apply to `loss_fn` are also applied
                  to `test_loss_fn`, if the latter is provided. Note that the
                  constraints on both `loss_fn` and `test_loss_fn` both depend on
                  `sample_wise_grads_per_batch`. This means `loss_fn` and
                  `test_loss_fn` must either both be "per-example"  loss functions,
                  or both be "reduction" loss functions. If not provided, the loss
                  function for test examples is assumed to be the same as the loss
                  function for training examples, i.e. `loss_fn`.
                  Default: None
          sample_wise_grads_per_batch (bool, optional): PyTorch's native gradient
                  computations w.r.t. model parameters aggregates the results for a
                  batch and does not allow to access sample-wise gradients w.r.t.
                  model parameters. This forces us to iterate over each sample in
                  the batch if we want sample-wise gradients which is computationally
                  inefficient. We offer an implementation of batch-wise gradient
                  computations w.r.t. to model parameters which is computationally
                  more efficient. This implementation can be enabled by setting the
                  `sample_wise_grad_per_batch` argument to `True`, and should be
                  enabled if and only if the `loss_fn` argument is a "reduction" loss
                  function. For example, `nn.BCELoss(reduction="sum")` would be a
                  valid `loss_fn` if this implementation is enabled (see
                  documentation for `loss_fn` for more details). Note that our
                  current implementation enables batch-wise gradient computations
                  only for a limited number of PyTorch nn.Modules: Conv2D and Linear.
                  This list will be expanded in the near future.  Therefore, please
                  do not enable this implementation if gradients will be computed
                  for other kinds of layers.
                  Default: False
      """
  
      self.model = model
  
      self.checkpoint = checkpoint
  
      self.checkpoints_load_func = checkpoints_load_func
      # actually load the checkpoint
      checkpoints_load_func(model, checkpoint)
      self.loss_fn = loss_fn
      # If test_loss_fn not provided, it's assumed to be same as loss_fn
      self.test_loss_fn = (
          loss_fn if test_loss_fn is None else test_loss_fn
      )
      self.sample_wise_grads_per_batch = sample_wise_grads_per_batch
      self.batch_size = batch_size
  
      if not isinstance(train_dataset, DataLoader):
          assert isinstance(batch_size, int), (
              "since the `train_dataset` argument was a `Dataset`, "
              "`batch_size` must be an int."
          )
          self.train_dataloader = DataLoader(
              train_dataset, batch_size, shuffle=False
          )
      else:
          self.train_dataloader = train_dataset
  
      if hessian_dataset is None:
          self.hessian_dataloader = self.train_dataloader
      elif not isinstance(hessian_dataset, DataLoader):
          assert isinstance(batch_size, int), (
              "since the `shared_dataset` argument was a `Dataset`, "
              "`batch_size` must be an int."
          )
          self.hessian_dataloader = DataLoader(
              hessian_dataset, batch_size, shuffle=False
          )
      else:
          self.hessian_dataloader = hessian_dataset
  
      # we check the loss functions in `InfluenceFunctionBase` rather than
      # individually in its child classes because we assume all its implementations
      # have the same requirements on loss functions, i.e. the type of reductions
      # supported. furthermore, these checks are done using a helper function that
      # handles all implementations with a `sample_wise_grads_per_batch`
      # initialization argument.
  
      # we save the reduction type for both `loss_fn` and `test_loss_fn` because
      # 1) if `sample_wise_grads_per_batch` is true, the reduction type is needed
      # to compute per-example gradients, and 2) regardless, reduction type for
      # `loss_fn` is needed to compute the Hessian.
  
      # check `loss_fn`
      self.reduction_type = _check_loss_fn(
          self, loss_fn, "loss_fn", sample_wise_grads_per_batch
      )
      # check `test_loss_fn` if it was provided
      if not (test_loss_fn is None):
          self.test_reduction_type = _check_loss_fn(
              self,
              test_loss_fn,
              "test_loss_fn",
              sample_wise_grads_per_batch,
          )
      else:
          self.test_reduction_type = self.reduction_type
  
      self.layer_modules = None
      if not (layers is None):
          self.layer_modules = _set_active_parameters(model, layers)
  
imports: import functools
  from abc import abstractmethod
  from operator import add
  from typing import Any, Callable, List, Optional, Tuple, Union
  import torch
  import torch.nn
  from captum._utils.gradient import _extract_parameters_from_layers
  from captum.influence._core.influence import DataInfluence
  from captum.influence._utils.common import _check_loss_fn, _compute_batch_loss_influence_function_base, _compute_jacobian_sample_wise_grads_per_batch, _dataset_fn, _flatten_params, _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, _params_to_names, _progress_bar_constructor, _self_influence_helper_intermediate_quantities_influence_function, _set_active_parameters, _top_eigen, _unflatten_params_factory, 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
label_desc: Use image node_id 1 to create a new InfluenceFunctionBase object from inherited base classes: DataInfluence with example: obj = InfluenceFunctionBase(model, train_dataset, checkpoint, checkpoints_load_func, layers, loss_fn, batch_size, hessian_dataset, test_loss_fn, sample_wise_grads_per_batch)