forked from ccf-ai-infra/GPUCodeForces
38 lines
1.1 KiB
Python
38 lines
1.1 KiB
Python
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 []
|