自定义Attention层必须继承tf.keras.layers.Layer，所有权重需在build中用add_weight声明，支持三阶张量输入、mask处理及MultiHead拼接后dense投影，确保可训练、可序列化、全模式兼容。
Attention层必须继承tf.keras.layers.Layer，不能只写函数
注意力计算逻辑如果只封装成普通Python函数（比如def scaled_dot_product_attention(...)），就无法被model.fit()追踪梯度、无法保存为SavedModel、也不能参与Layer级的权重管理。自定义Attention必须是可训练、可序列化的对象。
实操建议：
- 所有权重（如
W_q、W_k、W_v）必须在build(self, input_shape)中用self.add_weight()声明，不能在__init__里直接tf.Variable - 输入至少要支持
[batch, seq_len, features]三阶张量；若需兼容encoder-decoder场景，得额外接收mask参数并做tf.where掩码处理 - 别在
call()里用tf.print或print调试——Eager模式下看似有效，但Graph模式会静默失败
tf.linalg.band_part和tf.sequence_mask选哪个做attention mask
两者都能生成mask，但语义和适用场景完全不同：前者生成固定形状的带状矩阵（比如只保留上三角），后者根据每个样本的实际长度动态生成布尔掩码。decoder自回归mask要用tf.linalg.band_part，padding mask必须用tf.sequence_mask。
常见错误现象：InvalidArgumentError: Incompatible shapes，多因mask维度没对齐——比如用sequence_mask输出[batch, seq_len]，却直接和[batch, head, q_len, k_len]的attention logits相加。
实操建议：
- padding mask先扩展维度：
mask = mask[:, tf.newaxis, tf.newaxis, :]再广播 - causal mask用
tf.linalg.band_part(tf.ones((seq_len, seq_len)), -1, 0)，注意-1表示“下方无限”，0表示“上方0个” - 混合mask（如padding+causal）用
tf.minimum合并，别用tf.logical_and——返回bool类型，和float32的logits不兼容
MultiHeadAttention输出拼接后必须过self.dense，不能省略投影
原始论文和tf.keras.layers.MultiHeadAttention实现都要求：多个head的输出先拼接，再经一个线性变换（即WO矩阵）映射回原始维度。跳过这步会导致特征空间坍缩、梯度异常、下游层接收维度错乱。
性能影响明显：少这一层，模型收敛变慢，BLEU或准确率通常掉2–5个点，尤其在长序列上更敏感。
实操建议：
- 拼接后shape是
[batch, seq_len, num_heads * depth]，self.dense的units必须设为原始key_dim * num_heads的等效输出维（通常是key_dim * num_heads本身，或显式指定为output_dim） - 别复用
self.W_o做bias——Keras Layer要求bias由self.dense统一管理，否则model.save()时会漏掉 - 如果想复用TF官方层逻辑，直接调
tf.keras.layers.MultiHeadAttention的_compute_attention私有方法风险高，它不保证API稳定
自定义Attention层在tf.function下报TypeError: 'NoneType' object is not callable
这是最常踩的坑：在call()里用了未正确初始化的变量，或在build()前就访问了self.wq这类权重。Graph模式下，Keras不会像Eager那样延迟报错，而是在trace时崩溃。
根本原因是build()触发时机依赖于首次call()的input_shape；如果测试时传了None或空tensor，build()根本没执行，后续就调不到权重。
实操建议：
- 单元测试必须用真实shape调用一次：
layer(tf.random.normal((2, 10, 64)))，不能只layer.build(input_shape=(None, None, 64)) - 在
call()开头加assert hasattr(self, 'wq') and self.wq is not None，提前暴露问题 - 避免在
__init__里写self.wq = None再覆盖——Keras 2.10+对add_weight的惰性初始化更严格，直接删掉这类占位赋值
事情说清了就结束。真正难的不是写对公式，而是让Layer在SavedModel、TFLite、TPU all-reduce这些环节都不掉链子——每个add_weight、每个tf.where、每次维度扩展，都得同时扛住Eager调试和Graph部署两套约束。
