Make pairformer_emb_per_sample call into pairformer_emb

openfold3/core/model/heads/prediction_heads.py

@@ -145,15 +145,13 @@ def per_sample_pairformer_emb(
         )
 
         for i in range(no_samples):
-            zij_chunk = self.embed_zij(
-                si_input=si_input, zij=zij, x_pred=x_pred[:, i : i + 1]
-            )
-
-            si_chunk, zij_chunk = self.pairformer_stack(
-                si.clone(),  # Avoid inplace ops on si
-                zij_chunk,
-                single_mask,
-                pair_mask,
+            si_chunk, zij_chunk = self.pairformer_emb(
+                si_input=si_input,
+                si=si.clone(),  # Avoid inplace ops on si
+                zij=zij,
+                x_pred=x_pred[..., i : i + 1, :, :],
+                single_mask=single_mask,
+                pair_mask=pair_mask,
                 chunk_size=chunk_size,
                 use_deepspeed_evo_attention=use_deepspeed_evo_attention,
                 use_cueq_triangle_kernels=use_cueq_triangle_kernels,
@@ -302,6 +300,32 @@ def forward(
             zij:
                 [*, N_token, N_token, C_z] Updated pair representation
         """
+
+        batch_dim_counts = {
+            "x_pred": x_pred.dim() - 2,
+            "si_input": si_input.dim() - 2,
+            "si": si.dim() - 2,
+            "zij": zij.dim() - 3,
+            "single_mask": single_mask.dim() - 1,
+            "pair_mask": pair_mask.dim() - 2,
+        }
+
+        if len(set(batch_dim_counts.values())) != 1:
+            raise ValueError(
+                f"Inputs have different number of batch dimensions:\n"
+                f"{batch_dim_counts}.\n"
+                f"Shapes: x_pred={(*x_pred.shape,)},\n"
+                f"si_input={(*si_input.shape,)},\n"
+                f"si={(*si.shape,)}, zij={(*zij.shape,)},\n"
+                f"single_mask={(*single_mask.shape,)},\n"
+                f"pair_mask={(*pair_mask.shape,)}"
+            )
+
+        if apply_per_sample and batch_dim_counts["x_pred"] < 1:
+            raise ValueError(
+                "apply_per_sample is not compatible with no batch dimensions"
+            )
+
         if apply_per_sample:
             si, zij = self.per_sample_pairformer_emb(
                 si_input=si_input,
@@ -382,7 +406,9 @@ def _chunk(
         )
         no_samples = zij.shape[-4]
         for i in range(no_samples):
-            zij_out[:, i : i + 1] = self._compute_logits(zij[:, i : i + 1])
+            zij_out[..., i : i + 1, :, :, :] = self._compute_logits(
+                zij[..., i : i + 1, :, :, :]
+            )
 
         return zij_out
 
@@ -452,7 +478,9 @@ def _chunk(
         )
         no_samples = zij.shape[-4]
         for i in range(no_samples):
-            zij_out[:, i : i + 1] = self._compute_logits(zij[:, i : i + 1])
+            zij_out[..., i : i + 1, :, :, :] = self._compute_logits(
+                zij[..., i : i + 1, :, :, :]
+            )
 
         return zij_out
 

openfold3/tests/model_utils.py

@@ -0,0 +1,16 @@
+import torch
+
+from openfold3.core.model.primitives.initialization import lecun_normal_init_
+
+
+def initialize_model_weights(model):
+    """Re-initialize model weights with non-zero values.
+
+    Otherwise tests that check for agreement between different implementations
+    will effectively test nothing because the zero-initialized linear weights
+    will turn everything to zeros.
+    """
+    for module in model.modules():
+        if isinstance(module, torch.nn.Linear):
+            with torch.no_grad():
+                lecun_normal_init_(module.weight)

openfold3/tests/test_heads.py

@@ -32,6 +32,7 @@
 from openfold3.projects.of3_all_atom.project_entry import OF3ProjectEntry
 from openfold3.tests.config import consts
 from openfold3.tests.data_utils import random_of3_features
+from openfold3.tests.model_utils import initialize_model_weights
 
 
 class TestPredictedAlignedErrorHead(unittest.TestCase):
@@ -42,6 +43,7 @@ def test_predicted_aligned_error_head_shape(self):
         c_out = 50
 
         pae_head = PredictedAlignedErrorHead(c_z, c_out)
+        initialize_model_weights(pae_head)
 
         zij = torch.ones((batch_size, n_token, n_token, c_z))
         out = pae_head(zij)
@@ -58,6 +60,7 @@ def test_predicted_distance_error_head_shape(self):
         c_out = 50
 
         pde_head = PredictedDistanceErrorHead(c_z, c_out)
+        initialize_model_weights(pde_head)
 
         zij = torch.ones((batch_size, n_token, n_token, c_z))
         out = pde_head(zij)
@@ -76,6 +79,7 @@ def test_plddt_head_shape(self):
         plddt_head = PerResidueLDDTAllAtom(
             c_s, c_out, max_atoms_per_token=max_atoms_per_token.get()
         )
+        initialize_model_weights(plddt_head)
 
         si = torch.ones((batch_size, n_token, c_s))
         token_mask = torch.ones((batch_size, n_token))
@@ -107,6 +111,7 @@ def test_experimentally_resolved_head_all_atom_shape(self):
         exp_res_head = ExperimentallyResolvedHeadAllAtom(
             c_s, c_out, max_atoms_per_token=max_atoms_per_token.get()
         )
+        initialize_model_weights(exp_res_head)
 
         si = torch.ones((batch_size, n_token, c_s))
         token_mask = torch.ones((batch_size, n_token))
@@ -143,6 +148,7 @@ def test_pairformer_embedding_shape(self):
         pair_emb = PairformerEmbedding(
             **config.architecture.heads.pairformer_embedding
         ).eval()
+        initialize_model_weights(pair_emb)
 
         si_input = torch.ones(batch_size, n_token, c_s_input)
         si = torch.ones(batch_size, n_token, c_s)
@@ -174,6 +180,164 @@ def test_pairformer_embedding_shape(self):
         expected_shape_pair = (batch_size, n_token, n_token, c_z)
         np.testing.assert_array_equal(out_pair.shape, expected_shape_pair)
 
+    def test_pairformer_embedding_multisample_shape(self):
+        batch_size = consts.batch_size
+        n_sample = 5
+        n_token = consts.n_res
+
+        proj_entry = OF3ProjectEntry()
+        config = proj_entry.get_model_config_with_presets()
+
+        c_s_input = config.architecture.shared.c_s_input
+        c_s = config.architecture.shared.c_s
+        c_z = config.architecture.shared.c_z
+
+        pair_emb = PairformerEmbedding(
+            **config.architecture.heads.pairformer_embedding
+        ).eval()
+        initialize_model_weights(pair_emb)
+
+        si_input = torch.randn(batch_size, 1, n_token, c_s_input)
+        si = torch.randn(batch_size, 1, n_token, c_s)
+        zij = torch.randn(batch_size, 1, n_token, n_token, c_z)
+        atom_positions_predicted = torch.randn(batch_size, n_sample, n_token, 3)
+        single_mask = torch.randint(
+            0,
+            2,
+            size=(
+                batch_size,
+                1,
+                n_token,
+            ),
+        )
+        pair_mask = torch.randint(0, 2, size=(batch_size, 1, n_token, n_token))
+
+        out_single, out_pair = pair_emb(
+            si_input,
+            si,
+            zij,
+            atom_positions_predicted,
+            single_mask,
+            pair_mask,
+            chunk_size=4,
+        )
+
+        expected_shape_single = (batch_size, n_sample, n_token, c_s)
+        np.testing.assert_array_equal(out_single.shape, expected_shape_single)
+
+        expected_shape_pair = (batch_size, n_sample, n_token, n_token, c_z)
+        np.testing.assert_array_equal(out_pair.shape, expected_shape_pair)
+
+    def test_pairformer_embedding_apply_per_sample_equal(self):
+        batch_size = consts.batch_size
+        n_sample = 5
+        n_token = consts.n_res
+
+        proj_entry = OF3ProjectEntry()
+        config = proj_entry.get_model_config_with_presets()
+
+        c_s_input = config.architecture.shared.c_s_input
+        c_s = config.architecture.shared.c_s
+        c_z = config.architecture.shared.c_z
+
+        pair_emb = PairformerEmbedding(
+            **config.architecture.heads.pairformer_embedding
+        ).eval()
+        initialize_model_weights(pair_emb)
+
+        si_input = torch.randn(batch_size, 1, n_token, c_s_input)
+        si = torch.randn(batch_size, 1, n_token, c_s)
+        zij = torch.randn(batch_size, 1, n_token, n_token, c_z)
+        atom_positions_predicted = torch.randn(batch_size, n_sample, n_token, 3)
+        single_mask = torch.randint(
+            0,
+            2,
+            size=(
+                batch_size,
+                1,
+                n_token,
+            ),
+        )
+        pair_mask = torch.randint(0, 2, size=(batch_size, 1, n_token, n_token))
+
+        out_single, out_pair = pair_emb(
+            si_input,
+            si,
+            zij,
+            atom_positions_predicted,
+            single_mask,
+            pair_mask,
+            chunk_size=4,
+        )
+
+        out_single_per_sample, out_pair_per_sample = pair_emb(
+            si_input,
+            si,
+            zij,
+            atom_positions_predicted,
+            single_mask,
+            pair_mask,
+            chunk_size=4,
+            apply_per_sample=True,
+        )
+
+        torch.testing.assert_close(out_single, out_single_per_sample)
+        torch.testing.assert_close(out_pair, out_pair_per_sample)
+
+    def test_pairformer_embedding_apply_per_sample_one_batch_dim_equal(self):
+        n_sample = 5
+        n_token = consts.n_res
+
+        proj_entry = OF3ProjectEntry()
+        config = proj_entry.get_model_config_with_presets()
+
+        c_s_input = config.architecture.shared.c_s_input
+        c_s = config.architecture.shared.c_s
+        c_z = config.architecture.shared.c_z
+
+        pair_emb = PairformerEmbedding(
+            **config.architecture.heads.pairformer_embedding
+        ).eval()
+        initialize_model_weights(pair_emb)
+
+        si_input = torch.randn(1, n_token, c_s_input)
+        si = torch.randn(1, n_token, c_s)
+        zij = torch.randn(1, n_token, n_token, c_z)
+        atom_positions_predicted = torch.randn(n_sample, n_token, 3)
+        single_mask = torch.randint(
+            0,
+            2,
+            size=(
+                1,
+                n_token,
+            ),
+        )
+        pair_mask = torch.randint(0, 2, size=(1, n_token, n_token))
+
+        out_single, out_pair = pair_emb(
+            si_input,
+            si,
+            zij,
+            atom_positions_predicted,
+            single_mask,
+            pair_mask,
+            chunk_size=4,
+        )
+
+        out_single_per_sample, out_pair_per_sample = pair_emb(
+            si_input,
+            si,
+            zij,
+            atom_positions_predicted,
+            single_mask,
+            pair_mask,
+            chunk_size=4,
+            apply_per_sample=True,
+        )
+
+        torch.testing.assert_close(out_single, out_single_per_sample)
+        torch.testing.assert_close(out_pair, out_pair_per_sample)
+
 
 class TestAuxiliaryHeadsAllAtom(unittest.TestCase):
     def test_auxiliary_heads_all_atom_shape(self):
@@ -197,6 +361,7 @@ def test_auxiliary_heads_all_atom_shape(self):
         heads_config = config.architecture.heads
         heads_config.pae.enabled = True
         aux_head = AuxiliaryHeadsAllAtom(heads_config).eval()
+        initialize_model_weights(aux_head)
 
         si_input = torch.ones(batch_size, n_token, c_s_input)
         si = torch.ones(batch_size, n_token, c_s)