class_docstr: Returns a layer that maps `(Q, K, V, mask)` to `(activations, mask)`.
  
  This layer type performs the inner workings of one pass of multi-head
  self-attention. It:
  
    - subdivides incoming Q/K/V activations into multi-head versions;
    - for each head, computes the scaled dot product of each Q-K pair;
    - applies mask to screen out positions that come from padding tokens
      (indicated by 0 value);
    - [in ``'train'`` mode] applies dropout to Q-K dot products;
    - for each head, computes Q-K attention strengths using a per-query softmax
      of the Q-K dot products;
    - for each head, for each query position, combines V vectors according
      to the Q-K attention strengths; and
    - concatenates and fuses resulting per-head vectors into outgoing
      activations matching original input activation shapes.
class_docstr_tok: ["Returns","a","layer","that","maps","`","(","Q",",","K",",","V",",","mask",")","`","to","`","(","activations",",","mask",")","`",".","This","layer","type","performs","the","inner","workings","of","one","pass","of","multi-head","self-attention",".","It",":","-","subdivides","incoming","Q\/K\/V","activations","into","multi-head","versions",";","-","for","each","head",",","computes","the","scaled","dot","product","of","each","Q-K","pair",";","-","applies","mask","to","screen","out","positions","that","come","from","padding","tokens","(","indicated","by","0","value",")",";","-","[","in","``","'train","'","``","mode","]","applies","dropout","to","Q-K","dot","products",";","-","for","each","head",",","computes","Q-K","attention","strengths","using","a","per-query","softmax","of","the","Q-K","dot","products",";","-","for","each","head",",","for","each","query","position",",","combines","V","vectors","according","to","the","Q-K","attention","strengths",";","and","-","concatenates","and","fuses","resulting","per-head","vectors","into","outgoing","activations","matching","original","input","activation","shapes","."]
docstr: Returns attention-computed activations and unmodified mask.
  
  Args:
    inputs: A `(Q, K, V, mask)` tuple, whose query, key, and value
        activations have not yet been subdivided into heads.
code: def forward(self, inputs):
      """Returns attention-computed activations and unmodified mask.
  
      Args:
        inputs: A `(Q, K, V, mask)` tuple, whose query, key, and value
            activations have not yet been subdivided into heads.
      """
      q, k, v, mask = inputs
  
      d_feature = q.shape[-1]
      n_heads = self._n_heads
      if d_feature % n_heads != 0:
          raise ValueError(
              f"Dimensionality of feature embedding ({d_feature}) is not a "
              f"multiple of the requested number of attention heads ({n_heads})."
          )
  
      per_head_results, dots = _per_head_attention(
          SplitIntoHeads(n_heads, merged_batch_and_head=False).forward(
              q
          ),
          SplitIntoHeads(n_heads, merged_batch_and_head=False).forward(
              k
          ),
          SplitIntoHeads(n_heads, merged_batch_and_head=False).forward(
              v
          ),
          mask,
          dropout=self._dropout,
          mode=self._mode,
          rng=self.rng,
      )
      if self._mode == "viz":
          self.state = dots
      merged_results = MergeHeads(
          n_heads, merged_batch_and_head=False
      ).forward(per_head_results)
      return (merged_results, mask)
  
imports: import numpy
  from trax import fastmath
  from trax.fastmath import numpy
  from trax.layers import base
  from trax.layers import combinators
  from trax.layers import convolution
  from trax.layers import core
  from trax.layers import initializers
  from trax.layers.assert_shape import assert_shape
  from trax.layers.base import Fn
  from trax.layers.research import sparsity
label_desc: Use image node_id 2 for calling the PureAttention obj's underlying member method code with example usage: obj.forward(inputs) and returns: merged_results,  mask