class_docstr: Returns a layer that computes per-head attention (via scaled dot-product).
  
  This layer computes the core of the attention mechanism. Given per-head
  queries (Q), keys (K), values (V), and mask, it:
  
    - 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);
    - [if created in ``'train'`` mode] applies dropout to Q-K dot products;
    - computes Q-K attention strengths using a per-query softmax of the Q-K dot
      products; and
    - for each query position, combines V vectors according to the Q-K
      attention strengths.
class_docstr_tok: ["Returns","a","layer","that","computes","per-head","attention","(","via","scaled","dot-product",")",".","This","layer","computes","the","core","of","the","attention","mechanism",".","Given","per-head","queries","(","Q",")",",","keys","(","K",")",",","values","(","V",")",",","and","mask",",","it",":","-","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",")",";","-","[","if","created","in","``","'train","'","``","mode","]","applies","dropout","to","Q-K","dot","products",";","-","computes","Q-K","attention","strengths","using","a","per-query","softmax","of","the","Q-K","dot","products",";","and","-","for","each","query","position",",","combines","V","vectors","according","to","the","Q-K","attention","strengths","."]
docstr: Returns attention-computed per-head activations and unchanged mask.
  
  Args:
    inputs: A `(Q, K, V, mask)` tuple, whose query, key, and value
        activations have been subdivided into heads.
code: def forward(self, inputs):
      """Returns attention-computed per-head activations and unchanged mask.
  
      Args:
        inputs: A `(Q, K, V, mask)` tuple, whose query, key, and value
            activations have been subdivided into heads.
      """
      q, k, v, mask = inputs
      activations, attn_strengths = _per_head_attention(
          q,
          k,
          v,
          mask,
          dropout=self._dropout,
          mode=self._mode,
          rng=self.rng,
      )
      if self._mode == "viz":
          self.state = attn_strengths
      return activations
  
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 DotProductAttention obj's underlying member method code with example usage: obj.forward(inputs) and returns: activations