Corrected Attention docstring

Corrected Attention docstring missing head dimension for arguments mask and nonbatched_bias.
google-deepmind · Feb 5, 2024 · cb98745 · cb98745
1 parent 632ef57
commit cb98745
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/alphafold/model/modules.py b/alphafold/model/modules.py
@@ -558,8 +558,8 @@ def __call__(self, q_data, m_data, mask, nonbatched_bias=None):
       q_data: A tensor of queries, shape [batch_size, N_queries, q_channels].
       m_data: A tensor of memories from which the keys and values are
         projected, shape [batch_size, N_keys, m_channels].
-      mask: A mask for the attention, shape [batch_size, N_queries, N_keys].
-      nonbatched_bias: Shared bias, shape [N_queries, N_keys].
+      mask: A mask for the attention, shape [batch_size, N_heads, N_queries, N_keys].
+      nonbatched_bias: Shared bias, shape [N_heads, N_queries, N_keys].
 
     Returns:
       A float32 tensor of shape [batch_size, N_queries, output_dim].