code: def DotProductAttention(query, key, value, mask, dropout, mode, rng):
    
    depth = np.shape(query)[-1]
    dots = np.matmul(query, np.swapaxes(key, -1, -2)) / np.sqrt(depth)
    if mask is not None:
      dots = np.where(mask, dots, -1e9)
    # Softmax.
    dots = np.exp(dots - backend.logsumexp(dots, axis=-1, keepdims=True))
    if dropout >= 1.0:
      raise ValueError('Dropout rates must be lower than 1.')
    if dropout is not None and dropout > 0.0 and mode == 'train':
      keep = backend.random.bernoulli(rng, 1.0 - dropout, dots.shape)
      dots = np.where(keep, dots / (1.0 - dropout), 0)
    out = np.matmul(dots, value)
    return out
docs: Core dot product self-attention.
  
    Args:
      query: array of representations
      key: array of representations
      value: array of representations
      mask: attention-mask, gates attention
      dropout: float: dropout rate
      mode: 'eval' or 'train': whether to use dropout
      rng: JAX PRNGKey: subkey for disposable use
  
    Returns:
      Self attention for q, k, v arrays.