import torch
import torch.nn as nn
import math

class Model(nn.Module):
    """
    Log Beta operator implementation.
    """
    def __init__(self):
        super(Model, self).__init__()
    
    def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
        """
        Computes the logarithm of the beta function: log(beta(x, y))
        
        The beta function is defined as: beta(x, y) = gamma(x) * gamma(y) / gamma(x + y)
        So log(beta(x, y)) = lgamma(x) + lgamma(y) - lgamma(x + y)
        
        Args:
            x (torch.Tensor): First input tensor of any shape.
            y (torch.Tensor): Second input tensor of same shape as x.

        Returns:
            torch.Tensor: Output tensor with log beta applied, same shape as input.
        """
        return torch.lgamma(x) + torch.lgamma(y) - torch.lgamma(x + y)

batch_size = 512
dim = 16384

def get_inputs():
    x = torch.rand(batch_size, dim).cuda() * 10.0 + 0.1  # 避免零值
    y = torch.rand(batch_size, dim).cuda() * 10.0 + 0.1
    return [x, y]

def get_init_inputs():
    return []
