def PureMultiHeadedAttention(x, params, num_heads=8, dropout=0.0,
                             mode='train', **kwargs):
  """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).
  """
  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))
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).