code: def apply_mask(
      values: torch.FloatTensor,
      mask: Union[torch.BoolTensor, torch.IntTensor, torch.FloatTensor],
  ) -> torch.FloatTensor:
      """
      # Parameters
  
      values : `torch.FloatTensor`
          Shape `batch_size x num_attention_heads x source_seq_len x target_seq_len`
      mask : `torch.BoolTensor`
          Shape `batch_size x target_seq_len` OR `batch_size x 1 x 1 x target_seq_len`
      """
      # We create a 4D attention mask from a 2D or 3D tensor mask.
      if mask.dim() == 2:
          # The shape is `batch_size x 1 x 1 x target_seq_len` which is broadcast
          # to `batch_size x num_attention_heads x source_seq_len x target_seq_len`
          mask = mask[:, None, None, :]
      elif mask.dim() == 3:
          mask = mask[:, None, :, :]
      mask = mask.to(values.dtype)
      mask = (1.0 - mask) * min_value_of_dtype(values.dtype)
      return values + mask
  
imports: from typing import Union, Tuple
  import torch
  from allennlp.nn.util import min_value_of_dtype
label_desc: Use image node_id 1 for calling a global function with example usage: apply_mask(values, mask) and returns: unknown