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
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.
    
trax/layers/research/sparsity.py