Measures the average drop in model output as features are progressively removed in descending, ascending, and random importance order. Returns three scores (desc, asc, rand) plus optional baselines. ↑ desc, ↓ asc is preferred.

aopc

aopc(
    forward_func: Callable,
    inputs: TensorOrTupleOfTensorsGeneric,
    attributions: list[TensorOrTupleOfTensorsGeneric] | TensorOrTupleOfTensorsGeneric,
    baselines: BaselineType,
    feature_mask: TensorOrTupleOfTensorsGeneric = None,
    additional_forward_args: Any = None,
    target: ExplanationTarget | list[ExplanationTarget] = NoTarget(),
    frozen_features: list[Tensor] | None = None,
    max_features_processed_per_batch: int | None = None,
    total_feature_bins: int = 100,
    n_random_perms: int = 10,
    seed: int | None = None,
    multi_target: bool = False,
    show_progress: bool = True,
    return_intermediate_results: bool = False,
    return_dict: bool = False,
) -> Any

Implementation of Area over the Perturbation Curve by Samek et al., 2015. This implementation reuses the batch-computation ideas from captum and therefore it is fully compatible with the Captum library. In addition, the implementation takes some ideas about the implementation of the metric from the python Quantus library.

Consider a greedy iterative procedure that consists of measuring how the class encoded in the image (e.g. as measured by the function f) disappears when we progressively remove information from the image x, a process referred to as region perturbation, at the specified locations.

References

1) Wojciech Samek et al.: "Evaluating the visualization of what a deep neural network has learned." IEEE transactions on neural networks and learning systems 28.11 (2016): 2660-2673.

Parameters:

  • forward_func

    (Callable) –
    The forward function of the model or any modification of it.
    
  • inputs

    (Tensor or tuple[Tensor, ...]) –

    Input for which attributions are computed. If forward_func takes a single tensor as input, a single input tensor should be provided. If forward_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.

  • attributions

    (Tensor or tuple[Tensor, ...]) –
    Attribution scores computed based on an attribution algorithm.
    This attribution scores can be computed using the implementations
    provided in the `captum.attr` package. Some of those attribution
    approaches are so called global methods, which means that
    they factor in model inputs' multiplier, as described in:
    https://arxiv.org/abs/1711.06104
    Many global attribution algorithms can be used in local modes,
    meaning that the inputs multiplier isn't factored in the
    attribution scores.
    This can be done duing the definition of the attribution algorithm
    by passing `multipy_by_inputs=False` flag.
    For example in case of Integrated Gradients (IG) we can obtain
    local attribution scores if we define the constructor of IG as:
    ig = IntegratedGradients(multipy_by_inputs=False)
    
    Some attribution algorithms are inherently local.
    Examples of inherently local attribution methods include:
    Saliency, Guided GradCam, Guided Backprop and Deconvolution.
    
    For local attributions we can use real-valued perturbations
    whereas for global attributions that perturbation is binary.
    https://arxiv.org/abs/1901.09392
    
    If we want to compute the infidelity of global attributions we
    can use a binary perturbation matrix that will allow us to select
    a subset of features from `inputs` or `inputs - baselines` space.
    This will allow us to approximate sensitivity-n for a global
    attribution algorithm.
    
    Attributions have the same shape and dimensionality as the inputs.
    If inputs is a single tensor then the attributions is a single
    tensor as well. If inputs is provided as a tuple of tensors
    then attributions will be tuples of tensors as well.
    
  • baselines

    (scalar, Tensor, tuple of scalar, or Tensor) –
    Baselines define reference values against which the completeness is measured which sometimes
    represent ablated values and are used to compare with the actual inputs to compute
    importance scores in attribution algorithms. They can be represented
    as:
    
    - a single tensor, if inputs is a single tensor, with
      exactly the same dimensions as inputs or the first
      dimension is one and the remaining dimensions match
      with inputs.
    
    - a single scalar, if inputs is a single tensor, which will
      be broadcasted for each input value in input tensor.
    
    - a tuple of tensors or scalars, the baseline corresponding
      to each tensor in the inputs' tuple can be:
    
    - either a tensor with matching dimensions to
      corresponding tensor in the inputs' tuple
      or the first dimension is one and the remaining
      dimensions match with the corresponding
      input tensor.
    
    - or a scalar, corresponding to a tensor in the
      inputs' tuple. This scalar value is broadcasted
      for corresponding input tensor.
    
    Default: None
    
  • feature_mask

    (Tensor or tuple[Tensor, ...], default: None ) –
        feature_mask defines a mask for the input, grouping
        features which should be perturbed together. feature_mask
        should contain the same number of tensors as inputs.
        Each tensor should
        be the same size as the corresponding input or
        broadcastable to match the input tensor. Each tensor
        should contain integers in the range 0 to num_features
        - 1, and indices corresponding to the same feature should
        have the same value.
        Note that features within each input tensor are perturbed
        independently (not across tensors).
        If the forward function returns a single scalar per batch,
        we enforce that the first dimension of each mask must be 1,
        since attributions are returned batch-wise rather than per
        example, so the attributions must correspond to the
        same features (indices) in each input example.
        If None, then a feature mask is constructed which assigns
        each scalar within a tensor as a separate feature, which
        is perturbed independently.
        Default: None
    
  • additional_forward_args

    (Any, default: None ) –

    If the forward function requires additional arguments other than the inputs for which attributions should not be computed, this argument can be provided. It must be either a single additional argument of a Tensor or arbitrary (non-tuple) type or a tuple containing multiple additional arguments including tensors or any arbitrary python types. These arguments are provided to forward_func in order, following the arguments in inputs. Note that the perturbations are not computed with respect to these arguments.

    Default: None
    
  • target

    (int, tuple, Tensor, or list, default: NoTarget() ) –

    Indices for selecting predictions from output(for classification cases, this is usually the target class). If the network returns a scalar value per example, no target index is necessary. For general 2D outputs, targets can be either:

    - A single integer or a tensor containing a single
      integer, which is applied to all input examples
    
    - A list of integers or a 1D tensor, with length matching
      the number of examples in inputs (dim 0). Each integer
      is applied as the target for the corresponding example.
    
      For outputs with > 2 dimensions, targets can be either:
    
    - A single tuple, which contains #output_dims - 1
      elements. This target index is applied to all examples.
    
    - A list of tuples with length equal to the number of
      examples in inputs (dim 0), and each tuple containing
      #output_dims - 1 elements. Each tuple is applied as the
      target for the corresponding example.
    
    Default: None
    
  • max_features_processed_per_batch

    (int, default: None ) –

    The number of maximum input features that are processed together for every example. In case the number of features to be perturbed in each example (total_feature_bins) exceeds max_features_processed_per_batch, they will be sliced into batches of max_features_processed_per_batch examples and processed in a sequential order. However the total effective batch size will still be max_features_processed_per_batch * (2 + n_random_perms) as in each perturbation step, max_features_processed_per_batch * (2 + n_random_perms) features are processed. If max_features_processed_per_batch is None, all examples are processed together. max_features_processed_per_batch should at least be equal (2 + n_random_perms) and at most total_feature_bins * (2 + n_random_perms).

  • total_feature_bins

    (int, default: 100 ) –

    The total number of bins that the features will be perturbed in the descending, ascending and random order. Default: 100

  • frozen_features

    (List[Tensor], default: None ) –

    A list of frozen features that are not perturbed. This can be useful for ignoring the input structure features like padding, etc. Default: None In case CLS,PAD,SEP tokens are present in the input, they can be frozen by passing the indices of feature masks that correspond to these tokens.

  • n_random_perms

    (int, default: 10 ) –

    The number of random permutations of the feature importance scores that will be used to compute the AOPC scores for the random runs. Default: 10

  • seed

    (int, default: None ) –

    The seed value for the random number generator for reproducibility. Default: None

  • multi_target

    (bool, default: False ) –

    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. Default is False.

  • show_progress

    (bool, default: True ) –

    Displays the progress of the computation. Default: False

  • return_dict

    (bool, default: False ) –

    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.

Returns: A dictionary that contains the descending, ascending and random AOPC scores for the input samples. The dataclass contains the following fields: - desc (Union[List[Tensor], List[List[Tensor]]]): A list of tensors or a list of list of tensors representing the descending AOPC scores for each input example. The first dimension is equal to the number of examples in the input batch. - asc (Union[List[Tensor], List[List[Tensor]]]): A list of tensors or a list of list of tensors representing the ascending AOPC scores for each input example. The first dimension is equal to the number of examples in the input batch. - rand (Union[List[Tensor], List[List[Tensor]]]): A list of tensors or a list of list of tensors representing the random AOPC scores for each input example. The first dimension is equal to the number of examples in the input batch.

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) >>> baselines = torch.zeros(2, 3, 32, 32) >>> # Computes saliency maps for class 3. >>> attribution = saliency.attribute(input, target=3) >>> # define a perturbation function for the input

>>> # Computes the aopc scores for saliency maps
>>> aopc_desc, aopc_asc, aopc_rand = aopc(net, input, attribution, baselines)
Source code in torchxai/metrics/faithfulness/aopc.py
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
def aopc(
    forward_func: Callable,
    inputs: TensorOrTupleOfTensorsGeneric,
    attributions: list[TensorOrTupleOfTensorsGeneric] | TensorOrTupleOfTensorsGeneric,
    baselines: BaselineType,
    feature_mask: TensorOrTupleOfTensorsGeneric = None,
    additional_forward_args: Any = None,
    target: ExplanationTarget | list[ExplanationTarget] = NoTarget(),
    frozen_features: list[torch.Tensor] | None = None,
    max_features_processed_per_batch: int | None = None,
    total_feature_bins: int = 100,
    n_random_perms: int = 10,
    seed: int | None = None,
    multi_target: bool = False,
    show_progress: bool = True,
    return_intermediate_results: bool = False,
    return_dict: bool = False,
) -> Any:
    """
    Implementation of Area over the Perturbation Curve by Samek et al., 2015. This implementation
    reuses the batch-computation ideas from captum and therefore it is fully compatible with the Captum library.
    In addition, the implementation takes some ideas about the implementation of the metric from the python
    Quantus library.

    Consider a greedy iterative procedure that consists of measuring how the class
    encoded in the image (e.g. as measured by the function f) disappears when we
    progressively remove information from the image x, a process referred to as
    region perturbation, at the specified locations.

    References:
        1) Wojciech Samek et al.: "Evaluating the visualization of what a deep
        neural network has learned." IEEE transactions on neural networks and
        learning systems 28.11 (2016): 2660-2673.

    Args:
        forward_func (Callable):
                The forward function of the model or any modification of it.

        inputs (Tensor or tuple[Tensor, ...]): Input for which
                attributions are computed. If forward_func takes a single
                tensor as input, a single input tensor should be provided.
                If forward_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.

        attributions (Tensor or tuple[Tensor, ...]):
                Attribution scores computed based on an attribution algorithm.
                This attribution scores can be computed using the implementations
                provided in the `captum.attr` package. Some of those attribution
                approaches are so called global methods, which means that
                they factor in model inputs' multiplier, as described in:
                https://arxiv.org/abs/1711.06104
                Many global attribution algorithms can be used in local modes,
                meaning that the inputs multiplier isn't factored in the
                attribution scores.
                This can be done duing the definition of the attribution algorithm
                by passing `multipy_by_inputs=False` flag.
                For example in case of Integrated Gradients (IG) we can obtain
                local attribution scores if we define the constructor of IG as:
                ig = IntegratedGradients(multipy_by_inputs=False)

                Some attribution algorithms are inherently local.
                Examples of inherently local attribution methods include:
                Saliency, Guided GradCam, Guided Backprop and Deconvolution.

                For local attributions we can use real-valued perturbations
                whereas for global attributions that perturbation is binary.
                https://arxiv.org/abs/1901.09392

                If we want to compute the infidelity of global attributions we
                can use a binary perturbation matrix that will allow us to select
                a subset of features from `inputs` or `inputs - baselines` space.
                This will allow us to approximate sensitivity-n for a global
                attribution algorithm.

                Attributions have the same shape and dimensionality as the inputs.
                If inputs is a single tensor then the attributions is a single
                tensor as well. If inputs is provided as a tuple of tensors
                then attributions will be tuples of tensors as well.

        baselines (scalar, Tensor, tuple of scalar, or Tensor):
                Baselines define reference values against which the completeness is measured which sometimes
                represent ablated values and are used to compare with the actual inputs to compute
                importance scores in attribution algorithms. They can be represented
                as:

                - a single tensor, if inputs is a single tensor, with
                  exactly the same dimensions as inputs or the first
                  dimension is one and the remaining dimensions match
                  with inputs.

                - a single scalar, if inputs is a single tensor, which will
                  be broadcasted for each input value in input tensor.

                - a tuple of tensors or scalars, the baseline corresponding
                  to each tensor in the inputs' tuple can be:

                - either a tensor with matching dimensions to
                  corresponding tensor in the inputs' tuple
                  or the first dimension is one and the remaining
                  dimensions match with the corresponding
                  input tensor.

                - or a scalar, corresponding to a tensor in the
                  inputs' tuple. This scalar value is broadcasted
                  for corresponding input tensor.

                Default: None

        feature_mask (Tensor or tuple[Tensor, ...], optional):
                    feature_mask defines a mask for the input, grouping
                    features which should be perturbed together. feature_mask
                    should contain the same number of tensors as inputs.
                    Each tensor should
                    be the same size as the corresponding input or
                    broadcastable to match the input tensor. Each tensor
                    should contain integers in the range 0 to num_features
                    - 1, and indices corresponding to the same feature should
                    have the same value.
                    Note that features within each input tensor are perturbed
                    independently (not across tensors).
                    If the forward function returns a single scalar per batch,
                    we enforce that the first dimension of each mask must be 1,
                    since attributions are returned batch-wise rather than per
                    example, so the attributions must correspond to the
                    same features (indices) in each input example.
                    If None, then a feature mask is constructed which assigns
                    each scalar within a tensor as a separate feature, which
                    is perturbed independently.
                    Default: None

        additional_forward_args (Any, optional): If the forward function
                requires additional arguments other than the inputs for
                which attributions should not be computed, this argument
                can be provided. It must be either a single additional
                argument of a Tensor or arbitrary (non-tuple) type or a tuple
                containing multiple additional arguments including tensors
                or any arbitrary python types. These arguments are provided to
                forward_func in order, following the arguments in inputs.
                Note that the perturbations are not computed with respect
                to these arguments.

                Default: None
        target (int, tuple, Tensor, or list, optional): Indices for selecting
                predictions from output(for classification cases,
                this is usually the target class).
                If the network returns a scalar value per example, no target
                index is necessary.
                For general 2D outputs, targets can be either:

                - A single integer or a tensor containing a single
                  integer, which is applied to all input examples

                - A list of integers or a 1D tensor, with length matching
                  the number of examples in inputs (dim 0). Each integer
                  is applied as the target for the corresponding example.

                  For outputs with > 2 dimensions, targets can be either:

                - A single tuple, which contains #output_dims - 1
                  elements. This target index is applied to all examples.

                - A list of tuples with length equal to the number of
                  examples in inputs (dim 0), and each tuple containing
                  #output_dims - 1 elements. Each tuple is applied as the
                  target for the corresponding example.

                Default: None
        max_features_processed_per_batch (int, optional): The number of maximum input
                features that are processed together for every example. In case the number of
                features to be perturbed in each example (`total_feature_bins`) exceeds
                `max_features_processed_per_batch`, they will be sliced
                into batches of `max_features_processed_per_batch` examples and processed
                in a sequential order. However the total effective batch size will still be
                `max_features_processed_per_batch * (2 + n_random_perms)` as in each
                perturbation step, `max_features_processed_per_batch * (2 + n_random_perms)` features
                are processed. If `max_features_processed_per_batch` is None, all
                examples are processed together. `max_features_processed_per_batch` should
                at least be equal `(2 + n_random_perms)` and at most
                `total_feature_bins * (2 + n_random_perms)`.
        total_feature_bins (int, optional): The total number of bins that the features will be perturbed in
                the descending, ascending and random order. Default: 100
        frozen_features (List[torch.Tensor], optional): A list of frozen features that are not perturbed.
                This can be useful for ignoring the input structure features like padding, etc. Default: None
                In case CLS,PAD,SEP tokens are present in the input, they can be frozen by passing the indices
                of feature masks that correspond to these tokens.
        n_random_perms (int, optional): The number of random permutations of the feature importance scores
            that will be used to compute the AOPC scores for the random runs. Default: 10
        seed (int, optional): The seed value for the random number generator for reproducibility. 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.
                Default is False.
        show_progress (bool, optional): Displays the progress of the computation. Default: 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.
    Returns:
        A dictionary that contains the descending, ascending and random AOPC scores for the input
            samples. The dataclass contains the following fields:
            - desc (Union[List[Tensor], List[List[Tensor]]]): A list of tensors or a list of list of tensors
                representing the descending AOPC scores for each input example. The first dimension is equal to the
                number of examples in the input batch.
            - asc (Union[List[Tensor], List[List[Tensor]]]): A list of tensors or a list of list of tensors
                representing the ascending AOPC scores for each input example. The first dimension is equal to the
                number of examples in the input batch.
            - rand (Union[List[Tensor], List[List[Tensor]]]): A list of tensors or a list of list of tensors
                representing the random AOPC scores for each input example. The first dimension is equal to the
                number of examples in the input batch.

    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)
        >>> baselines = torch.zeros(2, 3, 32, 32)
        >>> # Computes saliency maps for class 3.
        >>> attribution = saliency.attribute(input, target=3)
        >>> # define a perturbation function for the input

        >>> # Computes the aopc scores for saliency maps
        >>> aopc_desc, aopc_asc, aopc_rand = aopc(net, input, attribution, baselines)
    """

    is_attributions_list = isinstance(attributions, list)
    is_targets_list = isinstance(target, list)
    if multi_target:
        assert is_targets_list, (
            "For multi-target explanations, target must be a list of targets"
        )
        target = [t.value for t in target]
        assert is_attributions_list, (
            "attributions must be a list of tensors or list of tuples of tensors"
        )
        assert is_targets_list, "targets must be a list of targets"
        assert all(isinstance(x, (tuple, int, list)) for x in target), (
            "targets must be a list of ints"
        )
        assert len(target) == len(
            attributions
        ), f"""The number of targets in the targets_list and
            attributions_list must match. Found number of targets in the targets_list is: {len(target)} and in the
            attributions_list: {len(attributions)}"""

    if not is_attributions_list:
        attributions = [attributions]
    if not is_targets_list:
        target = [target.value]  # type: ignore
    inputs_perturbed_aopc_desc_batch = []
    inputs_perturbed_aopc_asc_batch = []
    inputs_perturbed_aopc_rand_batch = []
    baselines_perturbed_aopc_desc_batch = []
    baselines_perturbed_aopc_asc_batch = []
    baselines_perturbed_aopc_rand_batch = []
    inputs_perturbed_fwds_agg_batch = []
    baselines_perturbed_fwds_agg_batch = []
    inputs_fwd_batch = []
    for a, t in tqdm.tqdm(zip(attributions, target, strict=False)):
        (
            inputs_perturbed_aopc_desc,
            inputs_perturbed_aopc_asc,
            inputs_perturbed_aopc_rand,
            baselines_perturbed_aopc_desc,
            baselines_perturbed_aopc_asc,
            baselines_perturbed_aopc_rand,
            inputs_perturbed_fwds_agg,
            baselines_perturbed_fwds_agg,
            inputs_fwd,
        ) = _aopc(
            forward_func=forward_func,
            inputs=inputs,
            attributions=a,
            baselines=baselines,
            feature_mask=feature_mask,
            additional_forward_args=additional_forward_args,
            target=t,  # type: ignore
            max_features_processed_per_batch=max_features_processed_per_batch,
            total_feature_bins=total_feature_bins,
            frozen_features=frozen_features,
            n_random_perms=n_random_perms,
            seed=seed,
            show_progress=show_progress,
        )
        inputs_perturbed_aopc_desc_batch.append(inputs_perturbed_aopc_desc)
        inputs_perturbed_aopc_asc_batch.append(inputs_perturbed_aopc_asc)
        inputs_perturbed_aopc_rand_batch.append(inputs_perturbed_aopc_rand)
        baselines_perturbed_aopc_desc_batch.append(baselines_perturbed_aopc_desc)
        baselines_perturbed_aopc_asc_batch.append(baselines_perturbed_aopc_asc)
        baselines_perturbed_aopc_rand_batch.append(baselines_perturbed_aopc_rand)
        inputs_perturbed_fwds_agg_batch.append(inputs_perturbed_fwds_agg)
        baselines_perturbed_fwds_agg_batch.append(baselines_perturbed_fwds_agg)
        inputs_fwd_batch.append(inputs_fwd)

    if not is_attributions_list:
        inputs_perturbed_aopc_desc_batch = inputs_perturbed_aopc_desc_batch[0]
        inputs_perturbed_aopc_asc_batch = inputs_perturbed_aopc_asc_batch[0]
        inputs_perturbed_aopc_rand_batch = inputs_perturbed_aopc_rand_batch[0]
        baselines_perturbed_aopc_desc_batch = baselines_perturbed_aopc_desc_batch[0]
        baselines_perturbed_aopc_asc_batch = baselines_perturbed_aopc_asc_batch[0]
        baselines_perturbed_aopc_rand_batch = baselines_perturbed_aopc_rand_batch[0]
        inputs_perturbed_fwds_agg_batch = inputs_perturbed_fwds_agg_batch[0]
        baselines_perturbed_fwds_agg_batch = baselines_perturbed_fwds_agg_batch[0]
        inputs_fwd_batch = inputs_fwd_batch[0]

    if return_intermediate_results:
        if return_dict:
            return {
                "desc": inputs_perturbed_aopc_desc_batch,
                "asc": inputs_perturbed_aopc_asc_batch,
                "rand": inputs_perturbed_aopc_rand_batch,
                "baselines_perturbed_desc": baselines_perturbed_aopc_desc_batch,
                "baselines_perturbed_asc": baselines_perturbed_aopc_asc_batch,
                "baselines_perturbed_rand": baselines_perturbed_aopc_rand_batch,
                "inputs_perturbed_fwds_agg_batch": inputs_perturbed_fwds_agg_batch,
                "baselines_perturbed_fwds_agg_batch": baselines_perturbed_fwds_agg_batch,
                "inputs_fwd_batch": inputs_fwd_batch,
            }
        else:
            return (
                inputs_perturbed_aopc_desc_batch,
                inputs_perturbed_aopc_asc_batch,
                inputs_perturbed_aopc_rand_batch,
                baselines_perturbed_aopc_desc_batch,
                baselines_perturbed_aopc_asc_batch,
                baselines_perturbed_aopc_rand_batch,
                inputs_perturbed_fwds_agg_batch,
                baselines_perturbed_fwds_agg_batch,
                inputs_fwd_batch,
            )
    else:
        if return_dict:
            return {
                "desc": inputs_perturbed_aopc_desc_batch,
                "asc": inputs_perturbed_aopc_asc_batch,
                "rand": inputs_perturbed_aopc_rand_batch,
                "baselines_perturbed_desc": baselines_perturbed_aopc_desc_batch,
                "baselines_perturbed_asc": baselines_perturbed_aopc_asc_batch,
                "baselines_perturbed_rand": baselines_perturbed_aopc_rand_batch,
            }
        else:
            return (
                inputs_perturbed_aopc_desc_batch,
                inputs_perturbed_aopc_asc_batch,
                inputs_perturbed_aopc_rand_batch,
                baselines_perturbed_aopc_desc_batch,
                baselines_perturbed_aopc_asc_batch,
                baselines_perturbed_aopc_rand_batch,
            )