code: def _convert_output_shape(
      unique_attr: Tensor,
      attr_inputs: Tuple[Tensor, ...],
      feature_mask: Tuple[Tensor, ...],
  ) -> Tuple[Tensor, ...]:
      """
      Convert the shape of a single tensor of unique feature attributionto
      to match the shape of the inputs returned by dataloader
      """
  
      # unique_attr in shape(*output_dims, n_features)
      output_dims = unique_attr.shape[:-1]
      n_features = unique_attr.shape[-1]
  
      attr = []
  
      for inp, mask in zip(attr_inputs, feature_mask):
          # input in shape(batch_size, *inp_feature_dims)
          # attribute in shape(*output_dims, *inp_feature_dims)
          attr_shape = (*output_dims, *inp.shape[1:])
  
          expanded_feature_indices = mask.expand(attr_shape)
  
          if len(inp.shape) > 2:
              # exclude batch_size & last of actual value
              extra_inp_dims = list(inp.shape[1:-1])
  
              # unsqueeze unqiue_attr to have same number of dims as inp
              # (*output_dims, 1..., 1, n_features)
              # then broadcast to (*output_dims, *inp.shape[1:-1], n_features)
              n_extra_dims = len(extra_inp_dims)
              unsqueezed_shape = (
                  *output_dims,
                  *(1,) * n_extra_dims,
                  n_features,
              )
              expanded_shape = (
                  *output_dims,
                  *extra_inp_dims,
                  n_features,
              )
              expanded_unqiue_attr = unique_attr.reshape(
                  unsqueezed_shape
              ).expand(expanded_shape)
          else:
              expanded_unqiue_attr = unique_attr
  
          # gather from (*output_dims, *inp.shape[1:-1], n_features)
          inp_attr = torch.gather(
              expanded_unqiue_attr, -1, expanded_feature_indices
          )
          attr.append(inp_attr)
  
      return tuple(attr)
  
file_path: captum/captum/attr/_core/dataloader_attr.py
imports: from collections import defaultdict
  from copy import copy
  from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple, Union
  import torch
  from captum._utils.common import _format_baseline, _format_feature_mask, _format_output, _format_tensor_into_tuples, _get_max_feature_index, _run_forward
  from captum._utils.typing import BaselineType
  from captum.attr._core.feature_ablation import FeatureAblation
  from captum.attr._utils.attribution import Attribution
  from torch import Tensor
label_desc: Use image node_id 4 for calling a global function with example usage: _convert_output_shape(unique_attr, attr_inputs, feature_mask) and returns: tuple