class_docstr: MultiHeadSelfAttention module optimized for Apple Neural Engine
      
class_docstr_tok: ["MultiHeadSelfAttention","module","optimized","for","Apple","Neural","Engine"]
docstr: Parameters:
      query: torch.tensor(bs, dim, 1, seq_length)
      key: torch.tensor(bs, dim, 1, seq_length)
      value: torch.tensor(bs, dim, 1, seq_length)
      mask: torch.tensor(bs, seq_length) or torch.tensor(bs, seq_length, 1, 1)
  
  Returns:
      weights: torch.tensor(bs, n_heads, seq_length, seq_length) Attention weights context: torch.tensor(bs,
      dim, 1, seq_length) Contextualized layer. Optional: only if `output_attentions=True`
docstr_tok: ["Parameters",":","query",":","torch.tensor","(","bs",",","dim",",","1",",","seq_length",")","key",":","torch.tensor","(","bs",",","dim",",","1",",","seq_length",")","value",":","torch.tensor","(","bs",",","dim",",","1",",","seq_length",")","mask",":","torch.tensor","(","bs",",","seq_length",")","or","torch.tensor","(","bs",",","seq_length",",","1",",","1",")","Returns",":","weights",":","torch.tensor","(","bs",",","n_heads",",","seq_length",",","seq_length",")","Attention","weights","context",":","torch.tensor","(","bs",",","dim",",","1",",","seq_length",")","Contextualized","layer",".","Optional",":","only","if","`","output_attentions=True","`"]
code: def forward(
      self,
      query,
      key,
      value,
      mask,
      head_mask=None,
      output_attentions=False,
  ):
      """
      Parameters:
          query: torch.tensor(bs, dim, 1, seq_length)
          key: torch.tensor(bs, dim, 1, seq_length)
          value: torch.tensor(bs, dim, 1, seq_length)
          mask: torch.tensor(bs, seq_length) or torch.tensor(bs, seq_length, 1, 1)
  
      Returns:
          weights: torch.tensor(bs, n_heads, seq_length, seq_length) Attention weights context: torch.tensor(bs,
          dim, 1, seq_length) Contextualized layer. Optional: only if `output_attentions=True`
      """
      # Parse tensor shapes for source and target sequences
      assert (
          len(query.size()) == 4
          and len(key.size()) == 4
          and len(value.size()) == 4
      )
  
      bs, dim, dummy, seqlen = query.size()
      # assert seqlen == key.size(3) and seqlen == value.size(3)
      # assert dim == self.dim
      # assert dummy == 1
  
      # Project q, k and v
      q = self.q_lin(query)
      k = self.k_lin(key)
      v = self.v_lin(value)
  
      # Validate mask
      if mask is not None:
          expected_mask_shape = [bs, seqlen, 1, 1]
          if mask.dtype == torch.bool:
              mask = mask.logical_not().float() * -1e4
          elif mask.dtype == torch.int64:
              mask = (1 - mask).float() * -1e4
          elif mask.dtype != torch.float32:
              raise TypeError(
                  f"Unexpected dtype for mask: {mask.dtype}"
              )
  
          if len(mask.size()) == 2:
              mask = mask.unsqueeze(2).unsqueeze(2)
  
          if list(mask.size()) != expected_mask_shape:
              raise RuntimeError(
                  f"Invalid shape for `mask` (Expected {expected_mask_shape}, got {list(mask.size())}"
              )
  
      if head_mask is not None:
          raise NotImplementedError
  
      # Compute scaled dot-product attention
      dim_per_head = self.dim // self.n_heads
      mh_q = q.split(
          dim_per_head, dim=1
      )  # (bs, dim_per_head, 1, max_seq_length) * n_heads
      mh_k = k.transpose(1, 3).split(
          dim_per_head, dim=3
      )  # (bs, max_seq_length, 1, dim_per_head) * n_heads
      mh_v = v.split(
          dim_per_head, dim=1
      )  # (bs, dim_per_head, 1, max_seq_length) * n_heads
  
      normalize_factor = float(dim_per_head) ** -0.5
      attn_weights = [
          torch.einsum("bchq,bkhc->bkhq", [qi, ki]) * normalize_factor
          for qi, ki in zip(mh_q, mh_k)
      ]  # (bs, max_seq_length, 1, max_seq_length) * n_heads
  
      if mask is not None:
          for head_idx in range(self.n_heads):
              attn_weights[head_idx] = attn_weights[head_idx] + mask
  
      attn_weights = [
          aw.softmax(dim=1) for aw in attn_weights
      ]  # (bs, max_seq_length, 1, max_seq_length) * n_heads
      attn = [
          torch.einsum("bkhq,bchk->bchq", wi, vi)
          for wi, vi in zip(attn_weights, mh_v)
      ]  # (bs, dim_per_head, 1, max_seq_length) * n_heads
  
      attn = torch.cat(attn, dim=1)  # (bs, dim, 1, max_seq_length)
  
      attn = self.out_lin(attn)
  
      if output_attentions:
          return attn, attn_weights.cat(dim=2)
      else:
          return (attn,)
  
imports: from ane_transformers.reference.layer_norm import LayerNormANE
  import torch
  import torch.nn
  from transformers.models.distilbert import modeling_distilbert
label_desc: Use image node_id 3 for calling the MultiHeadSelfAttention obj's underlying member method code with example usage: obj.forward(query, key, value, mask, head_mask, output_attentions) and returns: attn,  attn_weights, attn