(self, x, is_query)
return data_dash
Constructs nonnegative kernel features for fast softmax attention.

    Args:
      x: input for which features are computed.
      is_query: predicate indicating whether input data corresponds to
                queries or keys.

    Returns:
      Random features for fast softmax attention.
Constructs nonnegative kernel features for fast softmax attention.
Constructs nonnegative kernel features for fast softmax attention .
def nonnegative_softmax_kernel_feature_creator(self, x, is_query):
    """Constructs nonnegative kernel features for fast softmax attention.

    Args:
      x: input for which features are computed.
      is_query: predicate indicating whether input data corresponds to
                queries or keys.

    Returns:
      Random features for fast softmax attention.
    """
    if self._normalize_data:
      # We have e^{qk^T/sqrt{d}} = e^{q_norm k_norm^T}, where
      # w_norm = w * data_normalizer for w in {q,k}.
      data_normalizer = 1.0 / (jnp.sqrt(jnp.sqrt(x.shape[-1])))
    else:
      data_normalizer = 1.0
    ratio = 1.0 / jnp.sqrt(self._projection_matrix.shape[0])
    # TODO(wgaj): Double-check... Should there be only one batch dimension...?
    data_mod_shape = x.shape[0:1] + self._projection_matrix.shape
    data_thick_random_matrix = (jnp.zeros(data_mod_shape) +
                                self._projection_matrix)

    data_dash = jnp.einsum('Bij, Bkj -> Bik',
                           data_normalizer * x,
                           data_thick_random_matrix)
    diag_data = jnp.square(x)
    diag_data = jnp.sum(diag_data, axis=x.ndim - 1)
    diag_data = (diag_data / 2.0) * data_normalizer * data_normalizer
    diag_data = jnp.expand_dims(diag_data, axis=x.ndim - 1)

    last_dims_t = (len(data_dash.shape) - 1,)
    attention_dims_t = (1,)
    if is_query:
      data_dash = ratio * (
          jnp.exp(data_dash - diag_data -
                  jnp.max(data_dash, axis=last_dims_t, keepdims=True)) +
          self._epsilon)
    else:
      data_dash = ratio * (
          jnp.exp(data_dash - diag_data - jnp.max(
              data_dash, axis=last_dims_t + attention_dims_t, keepdims=True)) +
          self._epsilon)

    return data_dash
def nonnegative_softmax_kernel_feature_creator ( self , x , is_query ) : if self . _normalize_data : # We have e^{qk^T/sqrt{d}} = e^{q_norm k_norm^T}, where # w_norm = w * data_normalizer for w in {q,k}. data_normalizer = 1.0 / ( jnp . sqrt ( jnp . sqrt ( x . shape [ - 1 ] ) ) ) else : data_normalizer = 1.0 ratio = 1.0 / jnp . sqrt ( self . _projection_matrix . shape [ 0 ] ) # TODO(wgaj): Double-check... Should there be only one batch dimension...? data_mod_shape = x . shape [ 0 : 1 ] + self . _projection_matrix . shape data_thick_random_matrix = ( jnp . zeros ( data_mod_shape ) + self . _projection_matrix ) data_dash = jnp . einsum ( 'Bij, Bkj -> Bik' , data_normalizer * x , data_thick_random_matrix ) diag_data = jnp . square ( x ) diag_data = jnp . sum ( diag_data , axis = x . ndim - 1 ) diag_data = ( diag_data / 2.0 ) * data_normalizer * data_normalizer diag_data = jnp . expand_dims ( diag_data , axis = x . ndim - 1 ) last_dims_t = ( len ( data_dash . shape ) - 1 , ) attention_dims_t = ( 1 , ) if is_query : data_dash = ratio * ( jnp . exp ( data_dash - diag_data - jnp . max ( data_dash , axis = last_dims_t , keepdims = True ) ) + self . _epsilon ) else : data_dash = ratio * ( jnp . exp ( data_dash - diag_data - jnp . max ( data_dash , axis = last_dims_t + attention_dims_t , keepdims = True ) ) + self . _epsilon ) return data_dash