Add a note to MultiHeadAttention's docstring about calling _build_from_signature() when used in a custom Layer.
@@ -143,6 +143,12 @@ class MultiHeadAttention(Layer):
Finally, the result tensor with the last dimension as value_dim can take an
linear projection and return.
+ When using MultiHeadAttention inside a custom Layer, the custom Layer must
+ implement `build()` and call MultiHeadAttention's `_build_from_signature()`.
+ This enables weights to be restored correctly when the model is loaded.
+ TODO(b/172609172): link to documentation about calling custom build functions
+ when used in a custom Layer.
+
Examples:
Performs 1D cross-attention over two sequence inputs with an attention mask.
