import torch
import math


def relu_kernel_transformation(data,
                               is_query,
                               projection_matrix=None,
                               numerical_stabilizer=0.001):
    """Computes features for the ReLU-kernel.
    Computes random features for the ReLU kernel from
    https://arxiv.org/pdf/2009.14794.pdf.
    Args:
        data: input data tensor of the shape [B, L, H, D], where: B - batch
        dimension, L - attention dimensions, H - heads, D - features.
        is_query: indicates whether input data is a query oor key tensor.
        projection_matrix: random Gaussian matrix of shape [M, D], where M stands
        for the number of random features and each D x D sub-block has pairwise
        orthogonal rows.
        numerical_stabilizer: small positive constant for numerical stability.
    Returns:
        Corresponding kernel feature map.
    """
    del is_query
    if projection_matrix is None:
        return torch.nn.functional.relu(data) + numerical_stabilizer
    else:
        ratio = 1.0 / math.sqrt(torch.tensor(projection_matrix.shape[0]).float())
        data_dash = ratio * torch.einsum("blhd,md->blhm", data, projection_matrix)
        return torch.nn.functional.relu(data_dash) + numerical_stabilizer