code: def _scatter_itp_attr_by_mask(
      itp_attr: Tensor,
      input_shape: Tuple[int, ...],
      mask: Tensor,
  ) -> Tensor:
      """
      Scatter the attribution of the interpretable features to the model input shape
      by mask, if the interpretable features are the mask groups of the raw
      input elements,
      """
  
      # itp_attr in shape(*output_dims, n_itp_features)
      output_dims = itp_attr.shape[:-1]
      n_itp_features = itp_attr.shape[-1]
  
      # input_shape in shape(batch_size, *inp_feature_dims)
      # attribute in shape(*output_dims, *inp_feature_dims)
      attr_shape = (*output_dims, *input_shape[1:])
  
      expanded_feature_indices = mask.expand(attr_shape)
  
      if len(input_shape) > 2:
          # exclude batch_size & last of actual value
          extra_inp_dims = list(input_shape[1:-1])
  
          # unsqueeze itp_attr to have same number of dims as input
          # (*output_dims, 1..., 1, n_itp_features)
          # then broadcast to (*output_dims, *inp.shape[1:-1], n_itp_features)
          n_extra_dims = len(extra_inp_dims)
          unsqueezed_shape = (
              *output_dims,
              *(1,) * n_extra_dims,
              n_itp_features,
          )
          expanded_shape = (
              *output_dims,
              *extra_inp_dims,
              n_itp_features,
          )
          expanded_itp_attr = itp_attr.reshape(unsqueezed_shape).expand(
              expanded_shape
          )
      else:
          expanded_itp_attr = itp_attr
  
      # gather from (*output_dims, *inp.shape[1:-1], n_itp_features)
      attr = torch.gather(
          expanded_itp_attr, -1, expanded_feature_indices
      )
  
      return attr
  
file_path: captum/captum/attr/_utils/interpretable_input.py
imports: from typing import Any, Callable, Dict, List, Optional, Tuple, Union
  import torch
  from torch import Tensor
label_desc: Use image node_id 1 for calling a global function with example usage: _scatter_itp_attr_by_mask(itp_attr, input_shape, mask) and returns: attr