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 []