Clarify and fix documentation.

Addresses the following pull requests:
* #1083
* #1079
* #893

PiperOrigin-RevId: 798021344
Change-Id: I51bbc56998b99d53697b9d95e28cf1d55d562291

Author: rpachauri

alphafold/model/modules.py

@@ -556,10 +556,11 @@ def __call__(self, q_data, m_data, mask, nonbatched_bias=None):
 
     Arguments:
       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].
+      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_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].

run_alphafold.py

@@ -89,8 +89,9 @@ class ModelsToRelax(enum.Enum):
                     'database for use by JackHMMer.')
 flags.DEFINE_string('pdb70_database_path', None, 'Path to the PDB70 '
                     'database for use by HHsearch.')
-flags.DEFINE_string('pdb_seqres_database_path', None, 'Path to the PDB '
-                    'seqres database for use by hmmsearch.')
+flags.DEFINE_string('pdb_seqres_database_path', None, 'Full filepath to the '
+                    'PDB seqres database file (not just the directory) for use '
+                    'by hmmsearch.')
 flags.DEFINE_string('template_mmcif_dir', None, 'Path to a directory with '
                     'template mmCIF structures, each named <pdb_id>.cif')
 flags.DEFINE_string('max_template_date', None, 'Maximum template release date '