import torch
import torch.nn as nn

BATCH_SIZE = 4096
HIDDEN_DIM = 4096
SHAPE = (BATCH_SIZE, HIDDEN_DIM)

Q_VALUE = 0.9

class QuantileLoss(nn.Module):
    """
    Quantile Loss. L_q = q*|e| if e>=0 else (1-q)*|e|
    """
    def __init__(self, quantile=0.9, reduction='mean'):
        super(QuantileLoss, self).__init__()
        self.quantile = quantile
        self.reduction = reduction

    def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
        error = target - pred
        
        loss = torch.where(error >= 0, 
                         self.quantile * error, 
                         (1 - self.quantile) * -error)
        
        if self.reduction == 'mean':
            return loss.mean()
        elif self.reduction == 'sum':
            return loss.sum()
        return loss

class Model(nn.Module):
    def __init__(self, quantile=0.9, reduction='none'):
        super(Model, self).__init__()
        self.loss_fn = QuantileLoss(quantile=quantile, reduction=reduction)
    
    def forward(self, pred, target):
        return self.loss_fn(pred, target)

def get_inputs():
    pred = torch.randn(SHAPE, dtype=torch.float32)
    target = torch.randn(SHAPE, dtype=torch.float32)
    return [pred.contiguous(), target.contiguous()]

def get_init_inputs():
    return [Q_VALUE, 'none']