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
# 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`
allennlp/modules/transformer/util.py