forked from ccf-ai-infra/GPUCodeForces
25 lines
746 B
Python
25 lines
746 B
Python
import torch
|
|
import torch.nn as nn
|
|
|
|
class Model(nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
|
|
def forward(self, x: torch.Tensor, w: torch.Tensor) -> torch.Tensor:
|
|
# LogWeightedSumExp = log(sum(w * exp(x)))
|
|
# Stable implementation: max_x + log(sum(w * exp(x - max_x)))
|
|
max_x, _ = x.max(dim=-1, keepdim=True)
|
|
diff = x - max_x
|
|
sum_exp = torch.sum(w * torch.exp(diff), dim=-1)
|
|
return torch.log(sum_exp) + max_x.squeeze(-1)
|
|
|
|
batch_size = 1024
|
|
feature_dim = 4096
|
|
|
|
def get_inputs():
|
|
x = torch.randn(batch_size, feature_dim, dtype=torch.float32)
|
|
w = torch.rand(batch_size, feature_dim, dtype=torch.float32) # weights > 0
|
|
return [x, w]
|
|
|
|
def get_init_inputs():
|
|
return [] |