examples/research_projects/performer/modeling_flax_performer_utils.py/FastAttentionviaLowRankDecomposition/__init__
class FastAttentionviaLowRankDecomposition:    def __init__(
        self,
        matrix_creator,
        kernel_feature_creator,
        renormalize_attention,
        numerical_stabilizer,
        redraw_features,
        unidirectional,
        lax_scan_unroll=1,
    ):  # For optimal GPU performance, set to 16.
        rng = random.PRNGKey(0)
        self.matrix_creator = matrix_creator
        self.projection_matrix = self.draw_weights(rng)
        self.kernel_feature_creator = kernel_feature_creator
        self.renormalize_attention = renormalize_attention
        self.numerical_stabilizer = numerical_stabilizer
        self.redraw_features = redraw_features
        self.unidirectional = unidirectional
        self.lax_scan_unroll = lax_scan_unroll