code: def test_classification_baselines_as_function(self) -> None:
      num_in = 40
      inputs = torch.arange(0.0, num_in * 2.0).reshape(2, num_in)
  
      def generate_baselines() -> Tensor:
          return torch.arange(0.0, num_in * 4.0).reshape(4, num_in)
  
      def generate_baselines_with_inputs(inputs: Tensor) -> Tensor:
          inp_shape = cast(Tuple[int, ...], inputs.shape)
          return torch.arange(0.0, inp_shape[1] * 2.0).reshape(
              2, inp_shape[1]
          )
  
      def generate_baselines_returns_array() -> ndarray:
          return np.arange(0.0, num_in * 4.0).reshape(4, num_in)
  
      # 10-class classification model
      model = SoftmaxModel(num_in, 20, 10)
      model.eval()
      model.zero_grad()
  
      gradient_shap = GradientShap(model)
      n_samples = 10
      attributions, delta = gradient_shap.attribute(
          inputs,
          baselines=generate_baselines,
          target=torch.tensor(1),
          n_samples=n_samples,
          stdevs=0.009,
          return_convergence_delta=True,
      )
      _assert_attribution_delta(
          self, (inputs,), (attributions,), n_samples, delta
      )
  
      attributions, delta = gradient_shap.attribute(
          inputs,
          baselines=generate_baselines_with_inputs,
          target=torch.tensor(1),
          n_samples=n_samples,
          stdevs=0.00001,
          return_convergence_delta=True,
      )
      _assert_attribution_delta(
          self, (inputs,), (attributions,), n_samples, delta
      )
  
      with self.assertRaises(AssertionError):
          attributions, delta = gradient_shap.attribute(  # type: ignore
              inputs,
              # Intentionally passing wrong type.
              baselines=generate_baselines_returns_array,
              target=torch.tensor(1),
              n_samples=n_samples,
              stdevs=0.00001,
              return_convergence_delta=True,
          )
  
file_path: captum/tests/attr/test_gradient_shap.py
imports: from typing import cast, Tuple, Union
  import numpy
  import torch
  from captum._utils.typing import Tensor
  from captum.attr._core.gradient_shap import GradientShap
  from captum.attr._core.integrated_gradients import IntegratedGradients
  from numpy import ndarray
  from tests.helpers.basic import assertTensorAlmostEqual, BaseTest
  from tests.helpers.basic_models import BasicLinearModel, BasicModel2
  from tests.helpers.classification_models import SoftmaxModel
label_desc: Use image node_id 3 for calling the Test obj's underlying member method code with example usage: obj.test_classification_baselines_as_function() and returns: torch, torch, np