def generalized_kernel_feature_creator(
    data, projection_matrix, batch_dims_t, precision, kernel_fn, kernel_epsilon, normalize_data
):
    """
    Constructs kernel features for fast generalized attention

    Args:
      data: input for which features are computes
      projection_matrix: matrix used to compute features
      batch_dims_t: tuple of batch dimensions
      precision: precision parameter
      kernel_fn: kernel function used
      kernel_epsilon: additive positive term added to every feature for numerical
        stability
      normalize_data: predicate indicating whether data should be normalized

    Returns:
      Random features for fast generalized attention.
    """
    if normalize_data:
        data_normalizer = 1.0 / (jnp.sqrt(jnp.sqrt(data.shape[-1])))
    else:
        data_normalizer = 1.0
    if projection_matrix is None:
        return kernel_fn(data_normalizer * data) + kernel_epsilon
    else:
        data_mod_shape = data.shape[0 : len(batch_dims_t)] + projection_matrix.shape
        data_thick_random_matrix = jnp.zeros(data_mod_shape) + projection_matrix
        data_dash = lax.dot_general(
            data_normalizer * data,
            data_thick_random_matrix,
            (((data.ndim - 1,), (data_thick_random_matrix.ndim - 1,)), (batch_dims_t, batch_dims_t)),
            precision=precision,
        )
    data_prime = kernel_fn(data_dash) + kernel_epsilon
    return data_prime

    Constructs kernel features for fast generalized attention

    Args:
      data: input for which features are computes
      projection_matrix: matrix used to compute features
      batch_dims_t: tuple of batch dimensions
      precision: precision parameter
      kernel_fn: kernel function used
      kernel_epsilon: additive positive term added to every feature for numerical
        stability
      normalize_data: predicate indicating whether data should be normalized

    Returns:
      Random features for fast generalized attention.
    
dpkd/transformers/examples/research_projects/performer/modeling_flax_performer_utils.py