class_docstr: Multi-Head Attention optimized for efficient ANE deployment
      
docstr: Core routine for computing multi-head attention
  
  Args:
      q:              Projected query embeddings of shape (batch_size, d_qk, 1, tgt_seq_len)
      k:              Projected key embeddings of shape (batch_size, d_qk, 1, src_seq_len)
      v:              Projected value embeddings of shape (batch_size, d_v, 1, src_seq_len)
      qk_mask:        Float tensor of shape (batch_size, src_seq_len, 1, tgt_seq_len).
                      Indices with the a high negative value, e.g. -1e4, are excluded from attention
      k_mask:         Float tensor of shape (batch_size, src_seq_len, 1, 1).
                      Indices with the a high negative value, e.g. -1e4, are excluded from attention
  
  Returns:
      attn:           Attention embeddings of shape (batch_size, d_v, 1, tgt_seq_len)
      attn_weights:   If `return_weights` is True, returns the softmax attention weights used to compute the attention matrix
docstr_tok: ["Core","routine","for","computing","multi-head","attention","Args",":","q",":","Projected","query","embeddings","of","shape","(","batch_size",",","d_qk",",","1",",","tgt_seq_len",")","k",":","Projected","key","embeddings","of","shape","(","batch_size",",","d_qk",",","1",",","src_seq_len",")","v",":","Projected","value","embeddings","of","shape","(","batch_size",",","d_v",",","1",",","src_seq_len",")","qk_mask",":","Float","tensor","of","shape","(","batch_size",",","src_seq_len",",","1",",","tgt_seq_len",")",".","Indices","with","the","a","high","negative","value",",","e.g",".","-1e4",",","are","excluded","from","attention","k_mask",":","Float","tensor","of","shape","(","batch_size",",","src_seq_len",",","1",",","1",")",".","Indices","with","the","a","high","negative","value",",","e.g",".","-1e4",",","are","excluded","from","attention","Returns",":","attn",":","Attention","embeddings","of","shape","(","batch_size",",","d_v",",","1",",","tgt_seq_len",")","attn_weights",":","If","`","return_weights","`","is","True",",","returns","the","softmax","attention","weights","used","to","compute","the","attention","matrix"]
code: def _attention_fn(self, q, k, v, qk_mask, k_mask, return_weights):
      """Core routine for computing multi-head attention
  
      Args:
          q:              Projected query embeddings of shape (batch_size, d_qk, 1, tgt_seq_len)
          k:              Projected key embeddings of shape (batch_size, d_qk, 1, src_seq_len)
          v:              Projected value embeddings of shape (batch_size, d_v, 1, src_seq_len)
          qk_mask:        Float tensor of shape (batch_size, src_seq_len, 1, tgt_seq_len).
                          Indices with the a high negative value, e.g. -1e4, are excluded from attention
          k_mask:         Float tensor of shape (batch_size, src_seq_len, 1, 1).
                          Indices with the a high negative value, e.g. -1e4, are excluded from attention
  
      Returns:
          attn:           Attention embeddings of shape (batch_size, d_v, 1, tgt_seq_len)
          attn_weights:   If `return_weights` is True, returns the softmax attention weights used to compute the attention matrix
      """
      # Principle 2: Chunking Large Intermediate Tensors  (machinelearning.apple.com/research/apple-neural-engine)
      # Split q, k and v to compute a list of single-head attention functions
      mh_q = q.split(
          self.d_qk // self.n_head, dim=1
      )  # n_head * (batch_size, d_qk/n_head, 1, tgt_seq_len)
      # Principle 3: Minimizing Memory Copies
      # Avoid as many transposes and reshapes as possible
      mh_k = k.transpose(1, 3).split(
          self.d_qk // self.n_head, dim=3
      )  # n_head * (batch_size, src_seq_len, 1, d_qk/n_head)
      mh_v = v.split(
          self.d_v // self.n_head, dim=1
      )  # n_head * (batch_size, d_v/n_head, 1, src_seq_len)
  
      attn_weights = [
          torch.einsum("bchq,bkhc->bkhq", [qi, ki])
          * self.q_normalize_fact
          for qi, ki in zip(mh_q, mh_k)
      ]  # n_head * (batch_size, src_seq_len, 1, tgt_seq_len)
  
      # Apply attention masking
      if qk_mask is not None:
          for head_idx in range(self.n_head):
              attn_weights[head_idx] = attn_weights[head_idx] + qk_mask
      if k_mask is not None:
          for head_idx in range(self.n_head):
              attn_weights[head_idx] = attn_weights[head_idx] + k_mask
  
      attn_weights = [
          aw.softmax(dim=1) for aw in attn_weights
      ]  # n_head * (batch_size, src_seq_len, 1, tgt_seq_len)
      mh_w = [
          self.dropout(aw) for aw in attn_weights
      ]  # n_head * (batch_size, src_seq_len, 1, tgt_seq_len)
  
      attn = [
          torch.einsum("bkhq,bchk->bchq", wi, vi)
          for wi, vi in zip(mh_w, mh_v)
      ]  # n_head * (batch_size, d_v/n_head, 1, tgt_seq_len)
      attn = torch.cat(attn, dim=1)  # (batch_size, d_v, 1, tgt_seq_len)
  
      if return_weights:
          return attn, attn_weights
      return attn, None
  
imports: import torch
  import torch.nn
  from .layer_norm import LayerNormANE
label_desc: Use image node_id 3 for calling the MultiHeadAttention obj's underlying member method code with example usage: obj._attention_fn(q, k, v, qk_mask, k_mask, return_weights) and returns: attn,  None, attn,  attn_weights