How to use the captum.attr._utils.common._validate_input function in captum

def test_validate_input(self):
        with self.assertRaises(AssertionError):
            _validate_input(torch.Tensor([-1.0, 1.0]), torch.Tensor([-2.0]))
            _validate_input(
                torch.Tensor([-1.0, 1.0]), torch.Tensor([-1.0, 1.0]), n_steps=-1
            )
            _validate_input(
                torch.Tensor([-1.0, 1.0]), torch.Tensor([-1.0, 1.0]), method="abcde"
            )
        _validate_input(torch.Tensor([-1.0]), torch.Tensor([-2.0]))
        _validate_input(
            torch.Tensor([-1.0]), torch.Tensor([-2.0]), method="gausslegendre"
        )

def test_validate_input(self):
        with self.assertRaises(AssertionError):
            _validate_input(torch.Tensor([-1.0, 1.0]), torch.Tensor([-2.0]))
            _validate_input(
                torch.Tensor([-1.0, 1.0]), torch.Tensor([-1.0, 1.0]), n_steps=-1
            )
            _validate_input(
                torch.Tensor([-1.0, 1.0]), torch.Tensor([-1.0, 1.0]), method="abcde"
            )
        _validate_input(torch.Tensor([-1.0]), torch.Tensor([-2.0]))
        _validate_input(
            torch.Tensor([-1.0]), torch.Tensor([-2.0]), method="gausslegendre"
        )

def test_validate_input(self):
        with self.assertRaises(AssertionError):
            _validate_input(torch.Tensor([-1.0, 1.0]), torch.Tensor([-2.0]))
            _validate_input(
                torch.Tensor([-1.0, 1.0]), torch.Tensor([-1.0, 1.0]), n_steps=-1
            )
            _validate_input(
                torch.Tensor([-1.0, 1.0]), torch.Tensor([-1.0, 1.0]), method="abcde"
            )
        _validate_input(torch.Tensor([-1.0]), torch.Tensor([-2.0]))
        _validate_input(
            torch.Tensor([-1.0]), torch.Tensor([-2.0]), method="gausslegendre"
        )

def expand_and_update_baselines():
            def get_random_baseline_indices(bsz, baseline):
                num_ref_samples = baseline.shape[0]
                return np.random.choice(num_ref_samples, n_samples * bsz).tolist()

            # TODO allow to add noise to baselines as well
            # expand baselines to match the sizes of input
            if "baselines" not in kwargs:
                return

            baselines = kwargs["baselines"]
            baselines = _format_baseline(baselines, inputs)
            _validate_input(
                inputs, baselines, draw_baseline_from_distrib=draw_baseline_from_distrib
            )

            if draw_baseline_from_distrib:
                bsz = inputs[0].shape[0]
                baselines = tuple(
                    baseline[get_random_baseline_indices(bsz, baseline)]
                    if isinstance(baseline, torch.Tensor)
                    else baseline
                    for baseline in baselines
                )
            else:
                baselines = tuple(
                    baseline.repeat_interleave(n_samples, dim=0)
                    if isinstance(baseline, torch.Tensor)
                    and baseline.shape[0] == input.shape[0]

Examples::

                >>> # ImageClassifier takes a single input tensor of images Nx3x32x32,
                >>> # and returns an Nx10 tensor of class probabilities.
                >>> # It contains an attribute conv1, which is an instance of nn.conv2d,
                >>> # and the output of this layer has dimensions Nx12x32x32.
                >>> net = ImageClassifier()
                >>> lig = LayerIntegratedGradients(net, net.conv1)
                >>> input = torch.randn(2, 3, 32, 32, requires_grad=True)
                >>> # Computes layer integrated gradients for class 3.
                >>> # attribution size matches layer output, Nx12x32x32
                >>> attribution = lig.attribute(input, target=3)
        """
        inps, baselines = _format_input_baseline(inputs, baselines)
        _validate_input(inps, baselines, n_steps, method)

        baselines = _tensorize_baseline(inps, baselines)
        additional_forward_args = _format_additional_forward_args(
            additional_forward_args
        )

        if self.device_ids is None:
            self.device_ids = getattr(self.forward_func, "device_ids", None)
        inputs_layer = _forward_layer_eval(
            self.forward_func,
            inps,
            self.layer,
            device_ids=self.device_ids,
            additional_forward_args=additional_forward_args,
            attribute_to_layer_input=attribute_to_layer_input,
        )

Examples::

                >>> # ImageClassifier takes a single input tensor of images Nx3x32x32,
                >>> # and returns an Nx10 tensor of class probabilities.
                >>> # It contains an attribute conv1, which is an instance of nn.conv2d,
                >>> # and the output of this layer has dimensions Nx12x32x32.
                >>> net = ImageClassifier()
                >>> layer_int_inf = InternalInfluence(net, net.conv1)
                >>> input = torch.randn(2, 3, 32, 32, requires_grad=True)
                >>> # Computes layer internal influence.
                >>> # attribution size matches layer output, Nx12x32x32
                >>> attribution = layer_int_inf.attribute(input)
        """
        inputs, baselines = _format_input_baseline(inputs, baselines)
        _validate_input(inputs, baselines, n_steps, method)

        # Retrieve step size and scaling factor for specified approximation method
        step_sizes_func, alphas_func = approximation_parameters(method)
        step_sizes, alphas = step_sizes_func(n_steps), alphas_func(n_steps)

        # Compute scaled inputs from baseline to final input.
        scaled_features_tpl = tuple(
            torch.cat(
                [baseline + alpha * (input - baseline) for alpha in alphas], dim=0
            ).requires_grad_()
            for input, baseline in zip(inputs, baselines)
        )

        additional_forward_args = _format_additional_forward_args(
            additional_forward_args
        )

>>> # ImageClassifier takes a single input tensor of images Nx3x32x32,
            >>> # and returns an Nx10 tensor of class probabilities.
            >>> net = ImageClassifier()
            >>> ig = IntegratedGradients(net)
            >>> input = torch.randn(2, 3, 32, 32, requires_grad=True)
            >>> # Computes integrated gradients for class 3.
            >>> attribution = ig.attribute(input, target=3)
        """
        # Keeps track whether original input is a tuple or not before
        # converting it into a tuple.
        is_inputs_tuple = isinstance(inputs, tuple)

        inputs, baselines = _format_input_baseline(inputs, baselines)

        _validate_input(inputs, baselines, n_steps, method)

        # retrieve step size and scaling factor for specified approximation method
        step_sizes_func, alphas_func = approximation_parameters(method)
        step_sizes, alphas = step_sizes_func(n_steps), alphas_func(n_steps)

        # scale features and compute gradients. (batch size is abbreviated as bsz)
        # scaled_features' dim -> (bsz * #steps x inputs[0].shape[1:], ...)
        scaled_features_tpl = tuple(
            torch.cat(
                [baseline + alpha * (input - baseline) for alpha in alphas], dim=0
            ).requires_grad_()
            for input, baseline in zip(inputs, baselines)
        )

        additional_forward_args = _format_additional_forward_args(
            additional_forward_args

>>> # and the output of this layer has dimensions Nx12x32x32.
                >>> net = ImageClassifier()
                >>> neuron_cond = NeuronConductance(net, net.conv1)
                >>> input = torch.randn(2, 3, 32, 32, requires_grad=True)
                >>> # To compute neuron attribution, we need to provide the neuron
                >>> # index for which attribution is desired. Since the layer output
                >>> # is Nx12x32x32, we need a tuple in the form (0..11,0..31,0..31)
                >>> # which indexes a particular neuron in the layer output.
                >>> # Computes neuron conductance for neuron with
                >>> # index (4,1,2).
                >>> attribution = neuron_cond.attribute(input, (4,1,2))
        """
        is_inputs_tuple = isinstance(inputs, tuple)

        inputs, baselines = _format_input_baseline(inputs, baselines)
        _validate_input(inputs, baselines, n_steps, method)

        num_examples = inputs[0].shape[0]
        total_batch = num_examples * n_steps

        # Retrieve scaling factors for specified approximation method
        step_sizes_func, alphas_func = approximation_parameters(method)
        step_sizes, alphas = step_sizes_func(n_steps), alphas_func(n_steps)

        # Compute scaled inputs from baseline to final input.
        scaled_features_tpl = tuple(
            torch.cat(
                [baseline + alpha * (input - baseline) for alpha in alphas], dim=0
            ).requires_grad_()
            for input, baseline in zip(inputs, baselines)
        )

>>> # ImageClassifier takes a single input tensor of images Nx3x32x32,
            >>> # and returns an Nx10 tensor of class probabilities.
            >>> net = ImageClassifier()
            >>> # creates an instance of LayerDeepLift to interpret target
            >>> # class 1 with respect to conv4 layer.
            >>> dl = LayerDeepLift(net, net.conv4)
            >>> input = torch.randn(1, 3, 32, 32, requires_grad=True)
            >>> # Computes deeplift attribution scores for conv4 layer and class 3.
            >>> attribution = dl.attribute(input, target=1)
        """
        inputs = _format_input(inputs)
        baselines = _format_baseline(baselines, inputs)
        gradient_mask = apply_gradient_requirements(inputs)

        _validate_input(inputs, baselines)

        # set hooks for baselines
        self.model.apply(self._register_hooks_ref)

        baselines = _tensorize_baseline(inputs, baselines)

        attr_baselines = _forward_layer_eval(
            self.model,
            baselines,
            self.layer,
            additional_forward_args=additional_forward_args,
            attribute_to_layer_input=attribute_to_layer_input,
        )

        # remove forward hook set for baselines
        for forward_handles_ref in self.forward_handles_refs: