def Favor(d_feature, n_heads=1, n_random_features=256, dropout=0.0,
          numerical_stabilizer=0.001, use_approximate_softmax=False,
          scale_by_norm=0, normalize_data=False, epsilon=0.0001, mode='train'):
  """Returns a layer that maps (activations, mask) to (new_activations, mask).

  See the FAVOR paper for details: https://arxiv.org/abs/2006.03555

  Args:
    d_feature: Depth/dimensionality of feature embedding.
    n_heads: Number of attention heads.
    n_random_features: Free dimension size for the orthogonal random matrix.
    dropout: Probababilistic rate for internal dropout applied to attention
        activations (based on query-key pairs) before dotting them with values.
    numerical_stabilizer: float, small number used for numerical stability.
    use_approximate_softmax: Bool, if True uses approximate softmax, otherwise
                             Relu.
    scale_by_norm: Boolean; whether to scale orthogonal random matrix.
    normalize_data: predicate indicating whether data should be normalized.
    epsilon: numerical stabilizer.
    mode: One of `'train'`, `'eval'`, or `'predict'`.
  """
  del dropout  # not implemented yet but needed in the API

  return  tl.ConfigurableAttention(
      tl.Dense(d_feature), tl.Dense(d_feature), tl.Dense(d_feature),
      tl.Dense(d_feature),
      tl.FavorAttention(d_feature, n_heads, n_random_features,
                        numerical_stabilizer, use_approximate_softmax,
                        scale_by_norm, normalize_data, epsilon, mode),
      n_heads=n_heads)
Returns a layer that maps (activations, mask) to (new_activations, mask).

  See the FAVOR paper for details: https://arxiv.org/abs/2006.03555

  Args:
    d_feature: Depth/dimensionality of feature embedding.
    n_heads: Number of attention heads.
    n_random_features: Free dimension size for the orthogonal random matrix.
    dropout: Probababilistic rate for internal dropout applied to attention
        activations (based on query-key pairs) before dotting them with values.
    numerical_stabilizer: float, small number used for numerical stability.
    use_approximate_softmax: Bool, if True uses approximate softmax, otherwise
                             Relu.
    scale_by_norm: Boolean; whether to scale orthogonal random matrix.
    normalize_data: predicate indicating whether data should be normalized.
    epsilon: numerical stabilizer.
    mode: One of `'train'`, `'eval'`, or `'predict'`.
  
trax/layers/research/sparsity.py