examples/research_projects/performer/modeling_flax_performer_utils.py/make_fast_generalized_attention
def make_fast_generalized_attention(
    qkv_dim,
    renormalize_attention=True,
    numerical_stabilizer=0.0,
    nb_features=256,
    features_type="deterministic",
    kernel_fn=jax.nn.relu,
    kernel_epsilon=0.001,
    redraw_features=False,
    unidirectional=False,
    lax_scan_unroll=1,
):
    """Construct a fast generalized attention menthod."""
    logging.info("Fast generalized attention.: %s features and renormalize=%s", nb_features, renormalize_attention)
    if features_type == "ortho":
        matrix_creator = functools.partial(GaussianOrthogonalRandomMatrix, nb_features, qkv_dim, scaling=False)
    elif features_type == "iid":
        matrix_creator = functools.partial(GaussianUnstructuredRandomMatrix, nb_features, qkv_dim)
    elif features_type == "deterministic":
        matrix_creator = None
    else:
        raise ValueError("Unknown feature value type")

    def kernel_feature_creator(
        data, projection_matrix, attention_dims_t, batch_dims_t, precision, is_query, normalize_data=False
    ):
        del attention_dims_t
        del is_query
        return generalized_kernel_feature_creator(
            data, projection_matrix, batch_dims_t, precision, kernel_fn, kernel_epsilon, normalize_data
        )

    attention_fn = FastAttentionviaLowRankDecomposition(
        matrix_creator,
        kernel_feature_creator,
        renormalize_attention=renormalize_attention,
        numerical_stabilizer=numerical_stabilizer,
        redraw_features=redraw_features,
        unidirectional=unidirectional,
        lax_scan_unroll=lax_scan_unroll,
    ).dot_product_attention
    return attention_fn