如何使用Trax库SelfAttention类的多头参数，n_heads>1时报错如何解决？

问题原因 #

多头自注意力机制要求输入张量的最后一维（特征维度）必须可以被头数n_heads整除：自注意力层会将输入特征均匀拆分给每个注意力头单独计算，总特征数必须是头数的整数倍，否则会出现维度不匹配错误。
你当前的输入activations形状为(1, 100, 1)，最后一维特征数为1，设置n_heads=2时1无法被2整除，因此触发报错；n_heads=1时1可以被1整除，因此运行正常。

修复方法 #

调整输入的最后一维特征数，或者调整头数，保证满足输入最后一维 % n_heads == 0即可。
示例修改后可运行的代码如下：

import trax
import numpy as np

attention = trax.layers.SelfAttention(n_heads=2)

# 将输入最后一维调整为2，可被头数2整除
activations = np.random.randint(0, 10, (1, 100, 2)).astype(np.float32)
input = (activations, )

init = attention.init(input)
output = attention(input)

补充说明 #

实际使用中通常会将特征维度设置为16、32、128、512等数值，方便适配2、4、8、16等不同的头数配置。

内容的提问来源于stack exchange，提问作者Kenenbek Arzymatov