code: def attention(query, key, value, mask, d_k):
      """Scaled Dot Product Attention"""
      # mask shape: batch_size x 1 x seq_len x seq_len
  
      # scores shape: batch_size x num_heads x seq_len x seq_len
      scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(
          d_k
      )
      scores = scores.masked_fill(mask == 0, float("-inf"))
      # p_attn shape: batch_size x num_heads x seq_len x seq_len
      p_attn = F.softmax(scores, dim=3)
      # attn shape: batch_size x num_heads x seq_len x d_k
      attn = torch.matmul(p_attn, value)
      return attn, p_attn
  
imports: import copy
  import logging
  import math
  from enum import Enum
  import torch
  import torch.nn
  import torch.nn.functional
label_desc: Use image node_id 6 for calling a global function with example usage: attention(query, key, value, mask, d_k) and returns: attn,  p_attn