Coding Test 速查：Hot 100 + 手写 MHA / Conv

def backtrack(path, choices):
    if is_done(path):
        res.append(path.copy())
        return
    for choice in choices:
        if invalid(choice):
            continue
        path.append(choice)
        backtrack(path, next_choices)
        path.pop()

import math
import torch
import torch.nn as nn
import torch.nn.functional as F


class MultiHeadAttention(nn.Module):
    def __init__(self, embed_dim, num_heads, dropout=0.0, bias=True):
        super().__init__()
        assert embed_dim % num_heads == 0
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads

        self.q_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
        self.k_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
        self.v_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
        self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
        self.dropout = nn.Dropout(dropout)

    def _split_heads(self, x):
        # [B, T, C] -> [B, H, T, D]
        B, T, C = x.shape
        return x.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)

    def _merge_heads(self, x):
        # [B, H, T, D] -> [B, T, C]
        B, H, T, D = x.shape
        return x.transpose(1, 2).contiguous().view(B, T, H * D)

    def forward(self, x, mask=None, return_attn=False):
        # x: [B, T, C]
        q = self._split_heads(self.q_proj(x))
        k = self._split_heads(self.k_proj(x))
        v = self._split_heads(self.v_proj(x))

        scores = q @ k.transpose(-2, -1)
        scores = scores / math.sqrt(self.head_dim)

        if mask is not None:
            # bool mask: True means masked out; 0/1 mask: 0 means masked out
            if mask.dtype == torch.bool:
                scores = scores.masked_fill(mask, float("-inf"))
            else:
                scores = scores.masked_fill(mask == 0, float("-inf"))

        attn = F.softmax(scores, dim=-1)
        attn = self.dropout(attn)
        out = attn @ v
        out = self.out_proj(self._merge_heads(out))
        return (out, attn) if return_attn else out

import torch
import torch.nn.functional as F


def conv1d_naive(x, weight, bias=None, stride=1, padding=0, dilation=1):
    B, Cin, L = x.shape
    Cout, Cin_w, K = weight.shape
    assert Cin == Cin_w

    x = F.pad(x, (padding, padding))
    Lpad = x.shape[-1]
    Lout = (Lpad - dilation * (K - 1) - 1) // stride + 1
    out = x.new_zeros(B, Cout, Lout)

    for b in range(B):
        for co in range(Cout):
            for i in range(Lout):
                start = i * stride
                total = x.new_tensor(0.0)
                for ci in range(Cin):
                    for k in range(K):
                        total = total + x[b, ci, start + k * dilation] * weight[co, ci, k]
                if bias is not None:
                    total = total + bias[co]
                out[b, co, i] = total
    return out

torch.manual_seed(0)
x = torch.randn(2, 3, 12)
w = torch.randn(4, 3, 5)
b = torch.randn(4)
y = conv1d_naive(x, w, b, stride=2, padding=2)
ref = F.conv1d(x, w, b, stride=2, padding=2)
assert torch.allclose(y, ref, atol=1e-5)

import torch
import torch.nn.functional as F


def _pair(v):
    return v if isinstance(v, tuple) else (v, v)


def conv2d_naive(x, weight, bias=None, stride=1, padding=0, dilation=1):
    stride_h, stride_w = _pair(stride)
    pad_h, pad_w = _pair(padding)
    dil_h, dil_w = _pair(dilation)

    B, Cin, H, W = x.shape
    Cout, Cin_w, KH, KW = weight.shape
    assert Cin == Cin_w

    x = F.pad(x, (pad_w, pad_w, pad_h, pad_h))
    Hpad, Wpad = x.shape[-2], x.shape[-1]
    Hout = (Hpad - dil_h * (KH - 1) - 1) // stride_h + 1
    Wout = (Wpad - dil_w * (KW - 1) - 1) // stride_w + 1
    out = x.new_zeros(B, Cout, Hout, Wout)

    for b in range(B):
        for co in range(Cout):
            for oh in range(Hout):
                for ow in range(Wout):
                    h0 = oh * stride_h
                    w0 = ow * stride_w
                    total = x.new_tensor(0.0)
                    for ci in range(Cin):
                        for kh in range(KH):
                            for kw in range(KW):
                                total = total + (
                                    x[b, ci, h0 + kh * dil_h, w0 + kw * dil_w]
                                    * weight[co, ci, kh, kw]
                                )
                    if bias is not None:
                        total = total + bias[co]
                    out[b, co, oh, ow] = total
    return out

torch.manual_seed(0)
x = torch.randn(2, 3, 8, 9)
w = torch.randn(4, 3, 3, 5)
b = torch.randn(4)
y = conv2d_naive(x, w, b, stride=(2, 1), padding=(1, 2))
ref = F.conv2d(x, w, b, stride=(2, 1), padding=(1, 2))
assert torch.allclose(y, ref, atol=1e-5)