code: def PureMultiHeadedAttention(x, params, num_heads=8, dropout=0.0,
                               mode='train', **kwargs):
    
    del params
    rng = kwargs.get('rng', None)
    (q, k, v), mask = x
    feature_depth = q.shape[-1]
    assert feature_depth % num_heads == 0
    head_depth = feature_depth // num_heads
    nbatch = np.shape(q)[0]
    # nbatch, seqlen, feature_depth --> nbatch, num_heads, seqlen, head_depth
    def SplitHeads(x):
      return np.transpose(
          np.reshape(x, (nbatch, -1, num_heads, head_depth)), (0, 2, 1, 3))
    # nbatch, num_heads, seqlen, head_depth --> nbatch, seqlen, feature_depth
    def JoinHeads(x):  # pylint: disable=invalid-name
      return np.reshape(
          np.transpose(x, (0, 2, 1, 3)), (nbatch, -1, num_heads*head_depth))
    # Split heads, dot-product attention, rejoin heads.
    return JoinHeads(
        DotProductAttention(
            SplitHeads(q), SplitHeads(k), SplitHeads(v), mask,
            dropout=dropout, mode=mode, rng=rng))
docs: Pure transformer-style multi-headed attention.
  
    Args:
      x: inputs ((q, k, v), mask)
      params: parameters (none)
      num_heads: int: number of attention heads
      dropout: float: dropout rate
      mode: str: 'train' or 'eval'
      **kwargs: other arguments including the rng
  
    Returns:
      Pure Multi-headed attention layer (no Dense transforms on input).