class_docstr: Multi-Head Attention optimized for efficient ANE deployment
      
docstr: Args:
      q:                  Query embeddings of shape (batch_size, embed_dim, 1, tgt_seq_len)
      k:                  Key embeddings of shape (batch_size, embed_dim, 1, src_seq_len)
      v:                  Value embeddings of shape (batch_size, embed_dim, 1, src_seq_len)
      qpos:               Positional encodings for the query embeddings with same shape as `q`
      kpos:               Positional encodings for the key embeddings with same shape as `k`
      vpos:               Positional encodings for the key embeddings with same shape as `v`
      qk_mask:            Float tensor with shape (batch_size, src_seq_len, 1, tgt_seq_len). Example use case: for causal masking
                          in generative language models (e.g. GPT), fill the upper triangular part with a high negative value (e.g. -1e4).
                          Indices with the a high negative value, e.g. -1e4, are excluded from attention
      k_mask:             Float tensor with shape (batch_size, src_seq_len, 1, 1). Example use case: when excluding embeddings that
                          correspond to zero-padded pixels in an image or unused tokens in a text token sequence from attention.
                          Indices with the a high negative value, e.g. -1e4, are excluded from attention
      return_weights:     If True, returns the intermediate attention weights
  
  
  Note: If any of q,k,v has shape (batch_size, embed_dim, height, width) that represent a 2-d feature map, this will
  be flattened to (batch_size, embed_dim, 1, height * width)
  
  Note: `attn_weights` are never passed downstream even when return_weights=True because all the attn_weights
  are harvested from the outermost module (e.g. ane_transformers.model#Transformer) by means of forward hooks
docstr_tok: ["Args",":","q",":","Query","embeddings","of","shape","(","batch_size",",","embed_dim",",","1",",","tgt_seq_len",")","k",":","Key","embeddings","of","shape","(","batch_size",",","embed_dim",",","1",",","src_seq_len",")","v",":","Value","embeddings","of","shape","(","batch_size",",","embed_dim",",","1",",","src_seq_len",")","qpos",":","Positional","encodings","for","the","query","embeddings","with","same","shape","as","`","q","`","kpos",":","Positional","encodings","for","the","key","embeddings","with","same","shape","as","`","k","`","vpos",":","Positional","encodings","for","the","key","embeddings","with","same","shape","as","`","v","`","qk_mask",":","Float","tensor","with","shape","(","batch_size",",","src_seq_len",",","1",",","tgt_seq_len",")",".","Example","use","case",":","for","causal","masking","in","generative","language","models","(","e.g",".","GPT",")",",","fill","the","upper","triangular","part","with","a","high","negative","value","(","e.g",".","-1e4",")",".","Indices","with","the","a","high","negative","value",",","e.g",".","-1e4",",","are","excluded","from","attention","k_mask",":","Float","tensor","with","shape","(","batch_size",",","src_seq_len",",","1",",","1",")",".","Example","use","case",":","when","excluding","embeddings","that","correspond","to","zero-padded","pixels","in","an","image","or","unused","tokens","in","a","text","token","sequence","from","attention",".","Indices","with","the","a","high","negative","value",",","e.g",".","-1e4",",","are","excluded","from","attention","return_weights",":","If","True",",","returns","the","intermediate","attention","weights","Note",":","If","any","of","q",",","k",",","v","has","shape","(","batch_size",",","embed_dim",",","height",",","width",")","that","represent","a","2-d","feature","map",",","this","will","be","flattened","to","(","batch_size",",","embed_dim",",","1",",","height","*","width",")","Note",":","`","attn_weights","`","are","never","passed","downstream","even","when","return_weights=True","because","all","the","attn_weights","are","harvested","from","the","outermost","module","(","e.g",".","ane_transformers.model","#","Transformer",")","by","means","of","forward","hooks"]
code: def _forward_impl(
      self,
      q,
      k,
      v,
      qpos=None,
      kpos=None,
      vpos=None,
      qk_mask=None,
      k_mask=None,
      return_weights=True,
  ):
      """
      Args:
          q:                  Query embeddings of shape (batch_size, embed_dim, 1, tgt_seq_len)
          k:                  Key embeddings of shape (batch_size, embed_dim, 1, src_seq_len)
          v:                  Value embeddings of shape (batch_size, embed_dim, 1, src_seq_len)
          qpos:               Positional encodings for the query embeddings with same shape as `q`
          kpos:               Positional encodings for the key embeddings with same shape as `k`
          vpos:               Positional encodings for the key embeddings with same shape as `v`
          qk_mask:            Float tensor with shape (batch_size, src_seq_len, 1, tgt_seq_len). Example use case: for causal masking
                              in generative language models (e.g. GPT), fill the upper triangular part with a high negative value (e.g. -1e4).
                              Indices with the a high negative value, e.g. -1e4, are excluded from attention
          k_mask:             Float tensor with shape (batch_size, src_seq_len, 1, 1). Example use case: when excluding embeddings that
                              correspond to zero-padded pixels in an image or unused tokens in a text token sequence from attention.
                              Indices with the a high negative value, e.g. -1e4, are excluded from attention
          return_weights:     If True, returns the intermediate attention weights
  
  
      Note: If any of q,k,v has shape (batch_size, embed_dim, height, width) that represent a 2-d feature map, this will
      be flattened to (batch_size, embed_dim, 1, height * width)
  
      Note: `attn_weights` are never passed downstream even when return_weights=True because all the attn_weights
      are harvested from the outermost module (e.g. ane_transformers.model#Transformer) by means of forward hooks
      """
      # Parse tensor shapes for source and target sequences
      assert (
          len(q.size()) == 4
          and len(k.size()) == 4
          and len(v.size()) == 4
      )
      b, ct, ht, wt = q.size()
      b, cs, hs, ws = k.size()
  
      tgt_seq_len = ht * wt
      src_seq_len = hs * ws
  
      # Add positional encodings if any
      if qpos is not None:
          q = q + qpos
      if kpos is not None:
          k = k + kpos
      if vpos is not None:
          v = v + kpos
  
      # Project q,k,v
      q = self.q_proj(q)
      k = self.k_proj(k)
      v = self.v_proj(v)
  
      # Validate qk_mask (`attn_mask` in `torch.nn.MultiheadAttention`)
      expected_qk_mask_shape = [b, src_seq_len, 1, tgt_seq_len]
      if qk_mask is not None:
          if qk_mask.dtype != torch.float32:
              raise RuntimeError(
                  f"`qk_mask` must be of type torch.float32, received {qk_mask.dtype}"
              )
          if list(qk_mask.size()) != expected_qk_mask_shape:
              raise RuntimeError(
                  f"Invalid shape for `qk_mask` (Expected {expected_qk_mask_shape}, got {list(qk_mask.size())}"
              )
  
      # Validate k_mask (`key_padding_mask` in `torch.nn.MultiheadAttention`)
      expected_k_mask_shape = [b, src_seq_len, 1, 1]
      if k_mask is not None:
          if k_mask.dtype != torch.float32:
              raise RuntimeError(
                  f"`k_mask` must be of type torch.float32, received {k_mask.dtype}"
              )
          if list(k_mask.size()) != expected_k_mask_shape:
              raise RuntimeError(
                  f"Invalid shape for `k_mask` (Expected {expected_k_mask_shape}, got {list(k_mask.size())}"
              )
  
      # Call the attention function
      attn, attn_weights = self._attention_fn(
          q, k, v, qk_mask, k_mask, return_weights
      )
  
      # Revert to original dimension permutation
      attn = attn.contiguous().view(b, self.d_v, ht, wt)
  
      attn = self.out_proj(attn)
  
      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 4 for calling the MultiHeadAttention obj's underlying member method code with example usage: obj._forward_impl(q, k, v, qpos, kpos, vpos, qk_mask, k_mask, return_weights) and returns: attn,  None, attn,  attn_weights