import torch
import torch.nn as nn

class BatchNormModel(nn.Module):
    """
    Model that performs matrix multiplication followed by BatchNorm and ReLU activation.
    """
    def __init__(self, weight, num_features=2048, eps=1e-5, momentum=0.1, track_running_stats=True):
        super(BatchNormModel, self).__init__()
        self.weight = nn.Parameter(weight)
        # 设置 track_running_stats=True 以跟踪运行时统计量
        self.bn = nn.BatchNorm1d(
            num_features, 
            eps=eps, 
            momentum=momentum, 
            track_running_stats=track_running_stats
        )
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Performs matrix multiplication, applies BatchNorm, then ReLU activation.

        Args:
            x (torch.Tensor): Input tensor of shape [batch_size, input_dim]

        Returns:
            torch.Tensor: Output tensor of shape [batch_size, output_dim]
        """
        x = torch.matmul(x, self.weight)
        x = self.bn(x)
        return torch.relu(x)

# 添加别名以便在 run_code.py 中使用
Model = BatchNormModel

batch_size = 16
input_dim = 1024
output_dim = 2048

def get_inputs():
    x = torch.randn(batch_size, input_dim)
    return [x]

def get_init_inputs():
    weight = torch.randn(input_dim, output_dim)
    return [weight]