code: def form_attention_mask(input_mask, diagonal=None):
      """
      Build attention mask with optional masking of future tokens we forbid
      to attend to (e.g. as it is in Transformer decoder).
  
      Args:
          input_mask: binary mask of size B x L with 1s corresponding to valid
              tokens and 0s corresponding to padding tokens
          diagonal: diagonal where triangular future mask starts
              None -- do not mask anything
              0 -- regular translation or language modeling future masking
              1 -- query stream masking as in XLNet architecture
      Returns:
          attention_mask: mask of size B x 1 x L x L with 0s corresponding to
              tokens we plan to attend to and -10000 otherwise
      """
  
      if input_mask is None:
          return None
      attn_shape = (1, input_mask.shape[1], input_mask.shape[1])
      attn_mask = input_mask.to(dtype=bool).unsqueeze(1)
      if diagonal is not None:
          future_mask = torch.tril(
              torch.ones(
                  attn_shape, dtype=torch.bool, device=input_mask.device
              ),
              diagonal,
          )
          attn_mask = attn_mask & future_mask
      attention_mask = (1 - attn_mask.to(torch.float)) * NEG_INF
      return attention_mask.unsqueeze(1)
  
imports: import torch
  import torch.nn
label_desc: Use image node_id 1 for calling a global function with example usage: form_attention_mask(input_mask, diagonal) and returns: attention_mask, None