2025 USA-NA-AIO Round 1, Problem 2, Part 9

AI Olympiads 2025 USA-NA-AIO Round 1
1

Part 9 (5 points, coding task)

This question follows Part 6.

In this part, you are asked to program with PyTorch, not NumPy.

Build a deep neural network class named as Symmetric_Linear_Model that meets the modifications imposed in Part 6.

2 
### WRITE YOUR SOLUTION HERE ###

class Symmetric_Linear_Model(nn.Module):
    def __init__(self, in_features, hidden_features):
        super().__init__()
        self.in_features = in_features
        self.hidden_features = hidden_features
        self.linear0 = nn.Linear(in_features, hidden_features, bias = False)

    def forward(self, x):
        x = self.linear0(x)
        x = torch.sum(self.linear0.weight * x.reshape(*x.shape, 1), dim = -2)
        return x


""" END OF THIS PART """