code: def get_extended_attention_mask(
      attention_mask: torch.Tensor,
      input_shape: Tuple[int],
      causal_attention: bool = False,
  ) -> torch.Tensor:
      """
      Makes broadcastable attention and causal masks so that future and masked tokens are ignored.
      Arguments:
          attention_mask (`torch.Tensor`):
              Mask with ones indicating tokens to attend to, zeros for tokens to ignore.
          input_shape (`Tuple[int]`):
              The shape of the input to the model.
      Returns:
          `torch.Tensor` The extended attention mask, with a the same dtype as `attention_mask.dtype`.
  
      Method stolen from huggingface :)
      """
      # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
      # ourselves in which case we just need to make it broadcastable to all heads.
      if attention_mask.dim() == 3:
          extended_attention_mask = attention_mask[:, None, :, :]
      elif attention_mask.dim() == 2:
          # Provided a padding mask of dimensions [batch_size, seq_length]
          # - if the model is a decoder, apply a causal mask in addition to the padding mask
          # - if the model is an encoder, make the mask broadcastable to [batch_size, num_heads, seq_length, seq_length]
          if causal_attention:
              batch_size, seq_length = input_shape
              seq_ids = torch.arange(
                  seq_length, device=attention_mask.device
              )
              causal_mask = (
                  seq_ids[None, None, :].repeat(
                      batch_size, seq_length, 1
                  )
                  <= seq_ids[None, :, None]
              )
              # in case past_key_values are used we need to add a prefix ones mask to the causal mask
              # causal and attention masks must have same type with pytorch version < 1.3
              causal_mask = causal_mask.to(attention_mask.dtype)
  
              if causal_mask.shape[1] < attention_mask.shape[1]:
                  prefix_seq_len = (
                      attention_mask.shape[1] - causal_mask.shape[1]
                  )
                  causal_mask = torch.cat(
                      [
                          torch.ones(
                              (batch_size, seq_length, prefix_seq_len),
                              device=attention_mask.device,
                              dtype=causal_mask.dtype,
                          ),
                          causal_mask,
                      ],
                      axis=-1,
                  )
              extended_attention_mask = (
                  causal_mask[:, None, :, :]
                  * attention_mask[:, None, None, :]
              )
          else:
              extended_attention_mask = attention_mask[:, None, None, :]
      else:
          raise ValueError(
              f"Wrong shape for input_ids (shape {input_shape}) or attention_mask (shape {attention_mask.shape})"
          )
  
      # extended_attention_mask = extended_attention_mask.to(dtype=self.setup["dtype"])  # fp16 compatibility
      extended_attention_mask = (
          1.0 - extended_attention_mask
      ) * -10000.0
      return extended_attention_mask
  
imports: import torch
  import math
  from typing import Optional, Tuple
  from functools import partial
  from .embeddings import SinusoidalPositional, LearnablePositional, ScaledSinosoidal
  from .attention import get_attention_mechanism
label_desc: Use image node_id 3 for calling a global function with example usage: get_extended_attention_mask(attention_mask, input_shape, causal_attention) and returns: extended_attention_mask