Added checks for making sure the key and value depths are divisible by the number of attention heads when doing multihead attention.

PiperOrigin-RevId: 162025615

Author: T2T Team

tensor2tensor/models/common_attention.py

@@ -374,7 +374,18 @@ def multihead_attention(query_antecedent,
 
   Returns:
     A Tensor.
+
+  Raises:
+    ValueError: if the key depth or value depth are not divisible by the
+      number of attention heads.
   """
+  if total_key_depth % num_heads != 0:
+    raise ValueError("Key depth (%d) must be divisible by the number of "
+                     "attention heads (%d)." % (total_key_depth, num_heads))
+  if total_value_depth % num_heads != 0:
+    raise ValueError("Value depth (%d) must be divisible by the number of "
+                     "attention heads (%d)." % (total_value_depth, num_heads))
+
   with tf.variable_scope(
       name,
       default_name="multihead_attention",