examples/research_projects/performer/modeling_flax_performer_utils.py/kernel_feature_creator
def make_fast_softmax_attention(
    qkv_dim,
    renormalize_attention=True,
    numerical_stabilizer=0.000001,
    nb_features=256,
    ortho_features=True,
    ortho_scaling=0.0,
    redraw_features=True,
    unidirectional=False,
    nonnegative_features=True,
    lax_scan_unroll=1,
):
    """Construct a fast softmax attention method."""
    logging.info(
        "Fast softmax attention: %s features and orthogonal=%s, renormalize=%s",
        nb_features,
        ortho_features,
        renormalize_attention,
    )
    if ortho_features:
        matrix_creator = functools.partial(GaussianOrthogonalRandomMatrix, nb_features, qkv_dim, scaling=ortho_scaling)
    else:
        matrix_creator = functools.partial(GaussianUnstructuredRandomMatrix, nb_features, qkv_dim)
    if nonnegative_features:

        def kernel_feature_creator(
            data, projection_matrix, attention_dims_t, batch_dims_t, precision, is_query, normalize_data=True
        ):
            return nonnegative_softmax_kernel_feature_creator(
                data,
                projection_matrix,
                attention_dims_t,
                batch_dims_t,
                precision,
                is_query,
                normalize_data,
                numerical_stabilizer,
            )

    else:

        def kernel_feature_creator(
            data, projection_matrix, attention_dims_t, batch_dims_t, precision, is_query, normalize_data=True
        ):
            del is_query
            return sincos_softmax_kernel_feature_creator(
                data, projection_matrix, attention_dims_t, batch_dims_t, precision, 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