code: def _prepare_4d_attention_mask_for_sdpa(
      mask: torch.Tensor,
      dtype: torch.dtype,
      tgt_len: Optional[int] = None,
  ):
      """
      Creates a non-causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape
      `(batch_size, key_value_length)`
  
      Args:
          mask (`torch.Tensor` or `None`):
              A 2D attention mask of shape `(batch_size, key_value_length)`
          dtype (`torch.dtype`):
              The torch dtype the created mask shall have.
          tgt_len (`int`):
              The target length or query length the created mask shall have.
      """
      batch_size, key_value_length = mask.shape
      tgt_len = tgt_len if tgt_len is not None else key_value_length
  
      # torch.jit.trace and torchdynamo with fullgraph=True are unable to capture the controlflow `is_causal=attention_mask is None and q_len > 1`
      # used as an SDPA argument. We keep compatibility with these tracing tools by always using SDPA's `attn_mask` argument in case we are tracing.
      # TODO: Fix this as well when using torchdynamo with fullgraph=True.
      is_tracing = torch.jit.is_tracing()
  
      if torch.all(mask == 1):
          if is_tracing:
              pass
          elif tgt_len == 1:
              # For query_length == 1, causal attention and bi-directional attention are the same.
              return None
          elif key_value_length == tgt_len:
              return None
          else:
              # Unfortunately, for query_length > 1 and key_value_length != query_length, we can not generally ignore the attention mask, as SDPA causal mask generation
              # may be wrong. We will set is_causal=False in SDPA and rely on Transformers attention_mask instead, hence not setting it to None here.
              # Reference: https://github.com/pytorch/pytorch/issues/108108
              return AttentionMaskConverter._expand_mask(
                  mask=mask, dtype=dtype, tgt_len=tgt_len
              )
      else:
          return AttentionMaskConverter._expand_mask(
              mask=mask, dtype=dtype, tgt_len=tgt_len
          )
  
file_path: transformers/src/transformers/modeling_attn_mask_utils.py
imports: from dataclasses import dataclass
  from typing import List, Optional, Tuple, Union
  import torch
label_desc: Use image node_id 4 for calling a global function with example usage: _prepare_4d_attention_mask_for_sdpa(mask, dtype, tgt_len) and returns: AttentionMaskConverter, None, None, AttentionMaskConverter