掩码Softmax解析：给张量"戴面具"的魔法

想象你有一叠3D彩票券​（X），每张彩票有不同数量的有效号码​（valid_lens），而Softmax就像给这些号码按概率分配奖金。但有些号码是无效的​（比如彩票刮花了），我们得让它们不参与抽奖！

1. 什么是掩码Softmax？

普通Softmax：对所有数字按概率分配（比如 [1,2,3] → [0.09, 0.24, 0.67]）。

掩码Softmax：先给无效数字打马赛克​（换成 -1e6），再分配概率，让它们概率≈0！

2. 输入张量的形状变化（3D彩票券）​

假设输入 X 是一个 ​3D张量，形状是 [batch_size, num_rows, num_cols]：

• batch_size：有多少叠彩票（比如 2 个人各买一叠）。

• num_rows：每叠彩票有多少行（比如每叠 3 行）。

• num_cols：每行有多少个号码（比如每行 4 个号码）。

示例：

X = torch.randn(2, 3, 4)  # 2叠彩票，每叠3行，每行4个号码
valid_lens = torch.tensor([2, 3])  # 第1叠每行有效2个号码，第2叠每行有效3个号码

3. 掩码的"魔法步骤"

步骤 1：处理 valid_lens（统一成1D向量）

valid_lens 可以是两种形状：

1. ​1D向量​（比如 [2, 3]）：表示每叠彩票的所有行有效长度相同（第1叠每行有效2个号码，第2叠每行有效3个号码）。
2. ​2D向量​（比如 [[2,2,2], [3,3,3]]）：表示每叠彩票的每行有效长度可能不同（第1叠每行有效2个号码，第2叠每行有效3个号码）。

if valid_lens.dim() == 1:  # 如果是1D向量
    valid_lens = torch.repeat_interleave(valid_lens, shape[1])  # 复制成 [2,2,2,3,3,3]
else:  # 如果是2D向量
    valid_lens = valid_lens.reshape(-1)  # 展平成 [2,2,2,3,3,3]

结果：valid_lens 变成 ​1D向量，长度 = batch_size * num_rows（比如 6 个数字）。

步骤 2：给无效号码"打马赛克"（替换成 -1e6）​

把 X 先展平成 [batch_size * num_rows, num_cols]（比如 6 行 × 4 列），然后用 sequence_mask 在最后一列（dim=-1）​上：

• 如果号码有效 → 保留原值。
• 如果号码无效 → 替换成 -1e6（Softmax 后 ≈ 0）。

代码操作：

X = d2l.sequence_mask(X.reshape(-1, shape[-1]), valid_lens, value=-1e6)

效果：原本x的某些号码被替换为-1e6

[[1, 2, 3, 4],  
 [5, 6, 7, 8],  
 ...]  
→  
[[1, 2, -1e6, -1e6],  # 第1行只保留前2个号码  
 [5, 6, 7, -1e6],     # 第2行只保留前3个号码  
 ...]

步骤 3：对打过马赛克的张量做Softmax

现在 X 已经是 ​展平的形状 [batch_size * num_rows, num_cols]，直接对最后一列（dim=-1）做Softmax：

softmax_result = nn.functional.softmax(X, dim=-1)

结果： • 无效号码（-1e6）对应的Softmax概率 ≈ 0。 • 有效号码的概率正常分配（总和=1）。 最后再把 softmax_result ​还原成原始形状 [batch_size, num_rows, num_cols]：

return softmax_result.reshape(shape)

4. 总结

1. ​valid_lens 统一成1D：repeat_interleave 或 reshape，让所有彩票行的有效长度变一维数组。
2. ​打马赛克：无效号码 → -1e6（Softmax后≈0）。
3. ​Softmax抽奖：只对有效号码分配概率。
4. ​还原形状：变回原来的3D彩票堆！

幽默比喻：

• ​valid_lens = 彩票店的规则表（告诉每叠彩票哪些号码有效）。
• -1e6= 给无效号码泼墨水（让Softmax看不见它们）。
• ​Softmax = 公平分配奖金（无效号码拿不到钱！）。

这样，你的掩码Softmax就能正确处理带"刮痕"的彩票啦！