def sensitivity_max_and_avg(
    explainer: Explainer | Attribution,
    inputs: TensorOrTupleOfTensorsGeneric,
    perturb_func: Callable = default_perturb_func,
    perturb_radius: float = 0.02,
    n_perturb_samples: int = 10,
    norm_ord: str = "fro",
    max_examples_per_batch: int | None = None,
    multi_target: bool = False,
    return_dict: bool = False,
    **kwargs: Any,
) -> (
    Tensor
    | tuple[Tensor, Tensor]
    | list[Tensor]
    | tuple[list[Tensor], list[Tensor]]
    | dict[str, Tensor]
    | dict[str, list[Tensor]]
):
    r"""
    This is a modified version of the captum `captum.metric.sensitivity_max` (see: from captum.metrics import sensitivity_max)
    function that repeats the feature masks when performing perturbations. This function returns the sensitivity
    scores with both average and max aggregation resulting in sensitivity_max and sensitivity_avg together.

    The metric returns a list of senstivity scores if multi_target is True using the
    `torchxai.metrics.robustness.multi_target._multi_target_sensitivity_scores`,
    otherwise it returns a single sensitivity score same as the `captum.metric.sensitivity_max`.

    Explanation sensitivity measures the extent of explanation change when
    the input is slightly perturbed. It has been shown that the models that
    have high explanation sensitivity are prone to adversarial attacks:
    `Interpretation of Neural Networks is Fragile`
    https://www.aaai.org/ojs/index.php/AAAI/article/view/4252

    `sensitivity_max` metric measures maximum sensitivity of an explanation
    using Monte Carlo sampling-based approximation. By default in order to
    do so it samples multiple data points from a sub-space of an L-Infinity
    ball that has a `perturb_radius` radius using `default_perturb_func`
    default perturbation function. In a general case users can
    use any L_p ball or any other custom sampling technique that they
    prefer by providing a custom `perturb_func`.

    Note that max sensitivity is similar to Lipschitz Continuity metric
    however it is more robust and easier to estimate.
    Since the explanation, for instance an attribution function,
    may not always be continuous, can lead to unbounded
    Lipschitz continuity. Therefore the latter isn't always appropriate.

    More about the Lipschitz Continuity Metric can also be found here
    `On the Robustness of Interpretability Methods`
    https://arxiv.org/abs/1806.08049
    and
    `Towards Robust Interpretability with Self-Explaining Neural Networks`
    https://papers.nips.cc/paper\
    8003-towards-robust-interpretability-
    with-self-explaining-neural-networks.pdf

    More details about sensitivity max can be found here:
    `On the (In)fidelity and Sensitivity of Explanations`
    https://arxiv.org/abs/1901.09392

    Args:

        explanation_func (Callable):
                This function can be the `attribute` method of an
                attribution algorithm or any other explanation method
                that returns the explanations.

        inputs (Tensor or tuple[Tensor, ...]): Input for which
                explanations are computed. If `explanation_func` takes a
                single tensor as input, a single input tensor should
                be provided.
                If `explanation_func` takes multiple tensors as input, a tuple
                of the input tensors should be provided. It is assumed
                that for all given input tensors, dimension 0 corresponds
                to the number of examples (aka batch size), and if
                multiple input tensors are provided, the examples must
                be aligned appropriately.

        perturb_func (Callable):
                The perturbation function of model inputs. This function takes
                model inputs and optionally `perturb_radius` if
                the function takes more than one argument and returns
                perturbed inputs.

                If there are more than one inputs passed to sensitivity function those
                will be passed to `perturb_func` as tuples in the same order as they
                are passed to sensitivity function.

                It is important to note that for performance reasons `perturb_func`
                isn't called for each example individually but on a batch of
                input examples that are repeated `max_examples_per_batch / batch_size`
                times within the batch.

            Default: default_perturb_func
        perturb_radius (float, optional): The epsilon radius used for sampling.
            In the `default_perturb_func` it is used as the radius of
            the L-Infinity ball. In a general case it can serve as a radius of
            any L_p norm.
            This argument is passed to `perturb_func` if it takes more than
            one argument.

            Default: 0.02
        n_perturb_samples (int, optional): The number of times input tensors
                are perturbed. Each input example in the inputs tensor is
                expanded `n_perturb_samples` times before calling
                `perturb_func` function.

                Default: 10
        norm_ord (int, float, or str, optional): The type of norm that is used to
                compute the norm of the sensitivity matrix which is defined as the
                difference between the explanation function at its input and perturbed
                input. Acceptable values are either a string of 'fro' or 'nuc', or a
                number in the range of [-inf, inf] (including float("-inf") &
                float("inf")).

                Default: 'fro'
        max_examples_per_batch (int, optional): The number of maximum input
                examples that are processed together. In case the number of
                examples (`input batch size * n_perturb_samples`) exceeds
                `max_examples_per_batch`, they will be sliced
                into batches of `max_examples_per_batch` examples and processed
                in a sequential order. If `max_examples_per_batch` is None, all
                examples are processed together. `max_examples_per_batch` should
                at least be equal `input batch size` and at most
                `input batch size * n_perturb_samples`.

                Default: None
        multi_target (bool, optional): A boolean flag that indicates whether the metric computation is for
                multi-target explanations. if set to true, the targets are required to be a list of integers
                each corresponding to a required target class in the output. The corresponding metric outputs
                are then returned as a list of metric outputs corresponding to each target class. For multi-target
                infidelity, captum implementation is extened in _multi_target_infidelity function.
                Default is False.
        return_dict (bool, optional): A boolean flag that indicates whether the metric outputs are returned as a dictionary
                with keys as the metric names and values as the corresponding metric outputs.
                Default is False.
        **kwargs (Any, optional): Contains a list of arguments that are passed
                to `explanation_func` explanation function which in some cases
                could be the `attribute` function of an attribution algorithm.
                Any additional arguments that need be passed to the explanation
                function should be included here.
                For instance, such arguments include:
                `additional_forward_args`, `baselines` and `target`.

    Returns:

        sensitivities (Tensor): A tensor of scalar sensitivity scores per
                input example. The first dimension is equal to the
                number of examples in the input batch and the second
                dimension is one. Returned sensitivities are normalized by
                the magnitudes of the input explanations.

    Examples::
        >>> # ImageClassifier takes a single input tensor of images Nx3x32x32,
        >>> # and returns an Nx10 tensor of class probabilities.
        >>> net = ImageClassifier()
        >>> saliency = Saliency(net)
        >>> input = torch.randn(2, 3, 32, 32, requires_grad=True)
        >>> # Computes sensitivity score for saliency maps of class 3
        >>> sens = sensitivity_max(saliency.attribute, input, target = 3)
    """
    assert isinstance(explainer, (Explainer, Attribution)), (
        "The explainer must be an instance of the Explainer class or "
        "the Atribution class from the captum library."
    )
    if "explainer_baselines" in kwargs:
        kwargs["baselines"] = kwargs.pop("explainer_baselines")
    sensitivity_scores = _sensitivity_scores(
        explainer=explainer,
        inputs=inputs,
        perturb_func=perturb_func,
        perturb_radius=perturb_radius,
        n_perturb_samples=n_perturb_samples,
        norm_ord=norm_ord,
        max_examples_per_batch=max_examples_per_batch,
        multi_target=multi_target,
        **kwargs,
    )
    if multi_target:
        assert isinstance(sensitivity_scores, list), (
            "Expected sensitivity_scores to be a list of Tensors when multi_target is True."
        )
        sensitivity_max = [
            torch.max(score, dim=-1).values for score in sensitivity_scores
        ]
        sensitivity_avg = [torch.mean(score, dim=-1) for score in sensitivity_scores]
        if return_dict:
            return {
                "sensitivity_max": sensitivity_max,
                "sensitivity_avg": sensitivity_avg,
            }
        else:
            return sensitivity_max, sensitivity_avg
    else:
        assert isinstance(sensitivity_scores, Tensor), (
            "Expected sensitivity_scores to be a Tensor when multi_target is False."
        )
        sensitivity_max = torch.max(sensitivity_scores, dim=-1).values
        sensitivity_avg = torch.mean(sensitivity_scores, dim=-1)
        if return_dict:
            return {
                "sensitivity_max": sensitivity_max,
                "sensitivity_avg": sensitivity_avg,
            }
        else:
            return sensitivity_max, sensitivity_avg