code: def _format_feature_mask(
      feature_mask: Union[None, Tensor, Tuple[Tensor, ...]],
      inputs: Tuple[Tensor, ...],
  ) -> Tuple[Tensor, ...]:
      """
      Format a feature mask into a tuple of tensors.
      The `inputs` should be correctly formatted first
      If `feature_mask` is None, assign each non-batch dimension with a consecutive
      integer from 0.
      If `feature_mask` is a tensor, wrap it in a tuple.
      """
      if feature_mask is None:
          formatted_mask = []
          current_num_features = 0
          for inp in inputs:
              # the following can handle empty tensor where numel is 0
              # empty tensor will be added to the feature mask
              num_features = torch.numel(inp[0:1])
  
              formatted_mask.append(
                  current_num_features
                  + torch.reshape(
                      torch.arange(num_features, device=inp.device),
                      inp[0:1].shape,
                  )
              )
              current_num_features += num_features
          formatted_mask = tuple(formatted_mask)
  
      else:
          formatted_mask = _format_tensor_into_tuples(feature_mask)
  
      return formatted_mask
  
imports: import typing
  from enum import Enum
  from functools import reduce
  from inspect import signature
  from typing import Any, Callable, cast, Dict, List, overload, Tuple, Union
  import numpy
  import torch
  from captum._utils.typing import BaselineType, Literal, TargetType, TensorOrTupleOfTensorsGeneric, TupleOrTensorOrBoolGeneric
  from torch import device, Tensor
  from torch.nn import Module
label_desc: Use image node_id 11 for calling a global function with example usage: _format_feature_mask(feature_mask, inputs) and returns: formatted_mask