"""第 12 课补全练习：在归一化前遮挡未来。修改 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(scores):
    mask=torch.ones_like(scores,dtype=torch.bool).triu(1)
    return scores.masked_fill(mask,float('-inf')).softmax(-1)

w=solve(torch.zeros(3,3))
assert torch.equal(w.triu(1),torch.zeros(3,3))
torch.testing.assert_close(w.sum(-1),torch.ones(3))
assert w[0,0]==1
print("本课补全练习通过。")