code: def scaled_query_key_softmax(
      q: torch.Tensor,
      k: torch.Tensor,
      att_mask: Optional[
          Union[AttentionMask, "SparseCS", torch.Tensor]
      ],
  ) -> torch.Tensor:
      # TODO assume we have (N, S, hs) instead of (B, nh, S, hs), with N = B x nh
      # this is needed due to limitations in sparse_bmm for now
  
      # Self-attend: (N, S, hs) x (N, hs, S) -> (N, S, S)
      q = q / math.sqrt(k.size(-1))
  
      # Matmul with mask
      if att_mask is not None and isinstance(att_mask, AttentionMask):
          # Additive mask
          mask: Optional[
              Union[SparseCS, torch.Tensor]
          ] = att_mask.values
      else:
          mask = att_mask
  
      att = _matmul_with_mask(q, k.transpose(-2, -1), mask)
  
      # Softmax to get the attention probabilities
      is_causal = (
          isinstance(att_mask, AttentionMask) and att_mask.is_causal
      )
      att = _softmax(att, causal=is_causal)
      return att
  
file_path: xformers/xformers/components/attention/core.py
imports: import logging
  import math
  from contextlib import nullcontext
  from functools import lru_cache
  from typing import Optional, Union
  import torch
  from xformers import _has_cpp_library, _is_triton_available
  from xformers.components.attention.attention_mask import AttentionMask
label_desc: Use image node_id 7 for calling a global function with example usage: scaled_query_key_softmax(q, k, att_mask) and returns: att