code: def _naive_scaled_dot_product_attention(
      seqs: Tensor,
      keys: Tensor,
      key_padding_mask: Optional[PaddingMask],
      values: Tensor,
      attn_mask: Optional[AttentionMask],
      dropout_p: float,
      needs_weights: bool,
      training: bool,
  ) -> Tuple[Tensor, Optional[Tensor]]:
      # (N, H, S, K) @ (N, H, K, S_kv) = (N, H, S, S_kv)
      attn_weights = torch.matmul(seqs, keys.transpose(-1, -2))
  
      attn_weights = attn_weights * (seqs.size(-1) ** -0.5)
  
      if attn_mask is not None:
          # (S, S_kv)
          m = attn_mask.materialize()
  
          # (N, H, S, S_kv) + (S, S_kv) -> (N, H, S, S_kv)
          attn_weights = attn_weights + m
  
      if key_padding_mask is not None:
          # (N, S_kv)
          m = key_padding_mask.materialize()
  
          m = m[:, None, None, :]
  
          # (N, H, S, S_kv) + (N, 1, 1, S_kv) -> (N. H, S, S_kv)
          attn_weights = torch.where(m, attn_weights, -torch.inf)
  
      # For numerical stability run in single precision.
      attn_weights = softmax(attn_weights, dim=-1, dtype=torch.float32)
  
      attn_weights = attn_weights.type_as(seqs)
  
      if training and dropout_p > 0.0:
          attn_weights = dropout(attn_weights, dropout_p)
  
      # (N, H, S, S_kv) @ (N, H, S_kv, V) = (N, H, S, V)
      attn = torch.matmul(attn_weights, values)
  
      return attn, attn_weights if needs_weights else None
  
imports: import logging
  from abc import ABC, abstractmethod
  from contextlib import contextmanager
  from typing import Generator, Optional, Protocol, Tuple, final
  import torch
  import torch.nn.functional
  from torch import Tensor
  from torch.nn import Module
  from torch.nn.functional import dropout, softmax
  from fairseq2.nn.padding import PaddingMask
  from fairseq2.nn.transformer.attention_mask import AttentionMask, CausalAttentionMask
  from fairseq2.typing import finaloverride
label_desc: Use image node_id 1 for calling a global function with example usage: _naive_scaled_dot_product_attention(seqs, keys, key_padding_mask, values, attn_mask, dropout_p, needs_weights, training) and returns: attn,  unknown