"""第 13 课补全练习：把表示分成多个头。修改 solve，保持下方检查不变。"""
import math, io, argparse
import torch
from torch import nn
from torch.nn import functional as F
torch.set_num_threads(2)
torch.manual_seed(26)

def solve(x, heads):
    b,t,c=x.shape
    if c % heads: raise ValueError('宽度不能整除头数')
    return x.view(b,t,heads,c//heads).transpose(1,2)

x=torch.arange(24).view(1,3,8)
y=solve(x,2)
assert y.shape==(1,2,3,4)
assert torch.equal(y[0,1,0],x[0,0,4:])
print("本课补全练习通过。")