def DotProductAttention(query, key, value, mask, dropout, mode, rng):
  """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.
  """
  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
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 ) 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
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.