code: def PureDotProductAttention(dropout=0.0, mode='train'):
    
    def init_fun(_, input_shapes):  # pylint: disable=invalid-name
      q_shape, _, v_shape, _ = input_shapes
      output_shape = q_shape[:-1] + (v_shape[-1],)
      return output_shape, ()
    def apply_fun(params, inputs, **kwargs):  # pylint: disable=invalid-name
      del params
      q, k, v, mask = inputs
      rng = kwargs.get('rng', None)
      return DotProductAttention(q, k, v, mask,
                                 dropout=dropout, mode=mode, rng=rng)
    return init_fun, apply_fun
docs: Pure single-headed self-attention.
  
    Args:
      dropout: float: dropout rate
      mode: str: 'train' or 'eval'
  
    Returns:
      Pure single-headed attention layer. (No Dense transforms on input.)