code: def test_tracin_regression(
      self,
      reduction: Optional[str],
      tracin_constructor: Callable,
      mode: str,
      features: int,
      use_gpu: bool,
  ) -> None:
      with tempfile.TemporaryDirectory() as tmpdir:
          batch_size = 4
  
          dataset, net = self._test_tracin_regression_setup(
              tmpdir,
              features,
              use_gpu,
          )  # and not mode == 'sample_wise_trick'
  
          # check influence scores of training data
  
          train_inputs = dataset.samples
          train_labels = dataset.labels
  
          test_inputs = (
              torch.arange(17, 33, dtype=torch.float)
              .unsqueeze(1)
              .repeat(1, features)
          )
  
          if use_gpu:
              test_inputs = test_inputs.cuda()
  
          test_labels = test_inputs
  
          self.assertTrue(callable(tracin_constructor))
  
          if mode == "check_idx":
              self.assertTrue(isinstance(reduction, str))
              criterion = nn.MSELoss(reduction=cast(str, reduction))
  
              tracin = tracin_constructor(
                  net,
                  dataset,
                  tmpdir,
                  batch_size,
                  criterion,
              )
  
              train_scores = tracin.influence(
                  (train_inputs, train_labels)
              )
              idx, _ = tracin.influence(
                  (train_inputs, train_labels),
                  k=len(dataset),
                  proponents=True,
              )
              # check that top influence is one with maximal value
              # (and hence gradient)
              for i in range(len(idx)):
                  self.assertEqual(idx[i][0], 15)
  
              # check influence scores of test data
              test_scores = tracin.influence((test_inputs, test_labels))
              idx, _ = tracin.influence(
                  (test_inputs, test_labels),
                  k=len(test_inputs),
                  proponents=True,
              )
              # check that top influence is one with maximal value
              # (and hence gradient)
              for i in range(len(idx)):
                  self.assertTrue(_isSorted(idx[i]))
  
          if mode == "sample_wise_trick":
              criterion = nn.MSELoss(reduction="none")
  
              tracin = tracin_constructor(
                  net,
                  dataset,
                  tmpdir,
                  batch_size,
                  criterion,
                  sample_wise_grads_per_batch=False,
              )
  
              # With sample-wise trick
              criterion = nn.MSELoss(reduction="sum")
              tracin_sample_wise_trick = tracin_constructor(
                  net,
                  dataset,
                  tmpdir,
                  batch_size,
                  criterion,
                  sample_wise_grads_per_batch=True,
              )
  
              train_scores = tracin.influence(
                  (train_inputs, train_labels)
              )
              train_scores_sample_wise_trick = (
                  tracin_sample_wise_trick.influence(
                      (train_inputs, train_labels)
                  )
              )
              assertTensorAlmostEqual(
                  self, train_scores, train_scores_sample_wise_trick
              )
  
              test_scores = tracin.influence((test_inputs, test_labels))
              test_scores_sample_wise_trick = (
                  tracin_sample_wise_trick.influence(
                      (test_inputs, test_labels)
                  )
              )
              assertTensorAlmostEqual(
                  self, test_scores, test_scores_sample_wise_trick
              )
  
file_path: captum/tests/influence/_core/test_tracin_regression.py
imports: import os
  import tempfile
  from typing import Callable, cast, Optional
  import torch
  import torch.nn
  from captum.influence._core.arnoldi_influence_function import ArnoldiInfluenceFunction
  from captum.influence._core.influence_function import NaiveInfluenceFunction
  from captum.influence._core.tracincp import TracInCP
  from captum.influence._core.tracincp_fast_rand_proj import TracInCPFast, TracInCPFastRandProj
  from parameterized import parameterized
  from tests.helpers.basic import assertTensorAlmostEqual, BaseTest
  from tests.influence._utils.common import _isSorted, _wrap_model_in_dataparallel, build_test_name_func, CoefficientNet, DataInfluenceConstructor, IdentityDataset, RangeDataset
label_desc: Use image node_id 2 for calling the TestTracInRegression obj's underlying member method code with example usage: obj.test_tracin_regression(reduction, tracin_constructor, mode, features, use_gpu) without return types