Linear Layer - TorchCode

Problem Description

Implement a fully-connected linear layer: y = xW^T + b

Signature

class SimpleLinear:
    def __init__(self, in_features: int, out_features: int): ...
    def forward(self, x: torch.Tensor) -> torch.Tensor: ...

Requirements

• self.weight: shape (out_features, in_features), init with randn * (1/√in_features)

• self.bias: shape (out_features,), init as zeros

• Both must have requires_grad=True

• forward(x) computes x @ W^T + b

• Do NOT use torch.nn.Linear

Template

Implement the function below. Use only basic PyTorch operations.

# ✏️ YOUR IMPLEMENTATION HERE

class SimpleLinear:
    def __init__(self, in_features: int, out_features: int):
        pass  # Initialize weight and bias

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        pass  # Compute y = x @ W^T + b

Test Your Implementation

Use this code to debug before submitting.

# 🧪 Debug
layer = SimpleLinear(8, 4)
print("W shape:", layer.weight.shape)  # should be (4, 8)
print("b shape:", layer.bias.shape)    # should be (4,)

x = torch.randn(2, 8)
y = layer.forward(x)
print("Output shape:", y.shape)        # should be (2, 4)

Reference Solution

Try solving it yourself first! Click below to reveal the solution.

# ✅ SOLUTION

class SimpleLinear:
    def __init__(self, in_features: int, out_features: int):
        self.weight = torch.randn(out_features, in_features) * (1 / math.sqrt(in_features))
        self.weight.requires_grad_(True)
        self.bias = torch.zeros(out_features, requires_grad=True)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return x @ self.weight.T + self.bias

Tips

Run Locally

For interactive practice with auto-grading, run TorchCode locally:
pip install torch-judge then use check("linear")

Key Concepts

y = xW^T + b, Kaiming init, nn.Parameter