code: def add_outshape(module_masks, mask):
      """
      Inference the input mask of the add operation from the
      output mask.
      """
      assert isinstance(mask, CoarseMask)
  
      if module_masks.output_mask is None:
          module_masks.set_output_mask(mask)
          module_masks.set_input_mask(mask)
          return mask
      else:
          assert all(
              module_masks.output_mask.mask_index[1]
              == mask.mask_index[1]
          )
      return mask
  
imports: import logging
  import torch
label_desc: Use image node_id 6 for calling a global function with example usage: add_outshape(module_masks, mask) and returns: mask, mask