Torchcode Pytorch Interview Practice

from torch_judge import check, status, hint, reset_progress
# List all 40 problems and your progressstatus()
# Run tests for a specific problemcheck("relu")check("softmax")check("layernorm")check("attention")check("gpt2")
# Get a hint without spoilershint("softmax")
# Reset progress for a problemreset_progress("relu")

# Cell 1: Import judgefrom torch_judge import check, hintimport torchimport torch.nn as nn
# Cell 2: Your implementationdef my_relu(x: torch.Tensor) -> torch.Tensor:    # TODO: implement ReLU without using torch.relu or F.relu    raise NotImplementedError
# Cell 3: Run the judgecheck("relu")

def my_relu(x: torch.Tensor) -> torch.Tensor:    return torch.clamp(x, min=0)    # Alternative: return x * (x > 0)    # Alternative: return torch.where(x > 0, x, torch.zeros_like(x))

def my_softmax(x: torch.Tensor, dim: int = -1) -> torch.Tensor:    # Subtract max for numerical stability (prevents overflow)    x_max = x.max(dim=dim, keepdim=True).values    x_shifted = x - x_max    exp_x = torch.exp(x_shifted)    return exp_x / exp_x.sum(dim=dim, keepdim=True)

def my_layer_norm(    x: torch.Tensor,    weight: torch.Tensor,   # gamma (scale)    bias: torch.Tensor,     # beta (shift)    eps: float = 1e-5) -> torch.Tensor:    mean = x.mean(dim=-1, keepdim=True)    var = x.var(dim=-1, keepdim=True, unbiased=False)    x_norm = (x - mean) / torch.sqrt(var + eps)    return weight * x_norm + bias

def rms_norm(x: torch.Tensor, weight: torch.Tensor, eps: float = 1e-6) -> torch.Tensor:    rms = torch.sqrt((x ** 2).mean(dim=-1, keepdim=True) + eps)    return (x / rms) * weight

import torch.nn.functional as Fimport math
def scaled_dot_product_attention(    Q: torch.Tensor,  # (B, heads, T, head_dim)    K: torch.Tensor,    V: torch.Tensor,    mask: torch.Tensor = None) -> torch.Tensor:    d_k = Q.size(-1)    scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k)    if mask is not None:        scores = scores.masked_fill(mask == 0, float('-inf'))    attn_weights = F.softmax(scores, dim=-1)    return torch.matmul(attn_weights, V)

class MyMultiHeadAttention(nn.Module):    def __init__(self, d_model: int, num_heads: int):        super().__init__()        assert d_model % num_heads == 0        self.num_heads = num_heads        self.head_dim = d_model // num_heads        self.d_model = d_model
        self.W_q = nn.Linear(d_model, d_model)        self.W_k = nn.Linear(d_model, d_model)        self.W_v = nn.Linear(d_model, d_model)        self.W_o = nn.Linear(d_model, d_model)
    def forward(self, x: torch.Tensor, mask: torch.Tensor = None) -> torch.Tensor:        B, T, C = x.shape
        def split_heads(t):            return t.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
        Q = split_heads(self.W_q(x))        K = split_heads(self.W_k(x))        V = split_heads(self.W_v(x))
        attn_out = scaled_dot_product_attention(Q, K, V, mask)        # (B, heads, T, head_dim) -> (B, T, d_model)        attn_out = attn_out.transpose(1, 2).contiguous().view(B, T, C)        return self.W_o(attn_out)

def cross_entropy_loss(logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:    # logits: (B, C), targets: (B,) with class indices    # Use logsumexp trick for numerical stability    log_sum_exp = torch.logsumexp(logits, dim=-1)  # (B,)    log_probs = logits[torch.arange(len(targets)), targets]  # (B,)    return (log_sum_exp - log_probs).mean()

class MyDropout(nn.Module):    def __init__(self, p: float = 0.5):        super().__init__()        self.p = p
    def forward(self, x: torch.Tensor) -> torch.Tensor:        if not self.training or self.p == 0:            return x        mask = torch.bernoulli(torch.ones_like(x) * (1 - self.p))        return x * mask / (1 - self.p)  # inverted scaling

def kaiming_init(weight: torch.Tensor) -> torch.Tensor:    fan_in = weight.size(1)    std = math.sqrt(2.0 / fan_in)    with torch.no_grad():        weight.normal_(0, std)    return weight

def clip_grad_norm(parameters, max_norm: float) -> float:    params = [p for p in parameters if p.grad is not None]    total_norm = torch.sqrt(sum(p.grad.data.norm() ** 2 for p in params))    clip_coef = max_norm / (total_norm + 1e-6)    if clip_coef < 1:        for p in params:            p.grad.data.mul_(clip_coef)    return total_norm.item()

def train_with_accumulation(model, optimizer, dataloader, accumulation_steps=4):    optimizer.zero_grad()    for i, (inputs, targets) in enumerate(dataloader):        outputs = model(inputs)        loss = criterion(outputs, targets) / accumulation_steps  # scale loss        loss.backward()
        if (i + 1) % accumulation_steps == 0:            optimizer.step()            optimizer.zero_grad()

# WRONG — can overflow for large valuesexp_x = torch.exp(x)
# CORRECT — numerically stableexp_x = torch.exp(x - x.max(dim=-1, keepdim=True).values)

def causal_mask(T: int, device) -> torch.Tensor:    return torch.tril(torch.ones(T, T, device=device)).unsqueeze(0).unsqueeze(0)

class MyLayer(nn.Module):    def __init__(self, ...):        super().__init__()        self.weight = nn.Parameter(torch.empty(...))        self.bias = nn.Parameter(torch.zeros(...))        self._init_weights()
    def _init_weights(self):        nn.init.kaiming_uniform_(self.weight)
    def forward(self, x: torch.Tensor) -> torch.Tensor:        ...

def forward(self, x):    if self.training:        # use batch statistics        mean = x.mean(dim=0)        var = x.var(dim=0, unbiased=False)        # update running stats        self.running_mean = (1 - self.momentum) * self.running_mean + self.momentum * mean        self.running_var = (1 - self.momentum) * self.running_var + self.momentum * var    else:        # use running statistics        mean = self.running_mean        var = self.running_var    return (x - mean) / torch.sqrt(var + self.eps) * self.weight + self.bias

# WRONG — breaks autogradimport numpy as npresult = np.exp(x.numpy())
# CORRECT — autograd compatibleresult = torch.exp(x)