Skip to content

Make pairformer_emb_per_sample call into pairformer_emb #223

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
christinaflo merged 4 commits into from
May 20, 2026
Merged

Make pairformer_emb_per_sample call into pairformer_emb #223

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
4d9e242
Make pairformer_emb_per_sample call into pairformer_emb
GMNGeoffrey May 18, 2026
5af218c
Fix formatting
GMNGeoffrey May 18, 2026
a971fcc
Initialize model weights in test_heads.py
GMNGeoffrey May 19, 2026
505022f
Fix other assumed batch indices in prediction_heads.py
GMNGeoffrey May 19, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
42 changes: 33 additions & 9 deletions openfold3/core/model/heads/prediction_heads.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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, :, :],
Comment thread
GMNGeoffrey marked this conversation as resolved.
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,
Expand Down Expand Up @@ -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,
Expand Down
155 changes: 155 additions & 0 deletions openfold3/tests/test_heads.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -174,6 +174,161 @@ 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()

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(

@christinaflo christinaflo May 19, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This should be a separate utility, but for asserting equal tensors you have to modify the model weight init scheme like here: https://github.com/aqlaboratory/openfold-3/blob/main/openfold3/tests/test_kernels.py#L411
Otherwise, there's a lot of "final" init layers that are zero weight/bias on the first pass and you won't see the diffs come up

@GMNGeoffrey GMNGeoffrey May 19, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Oh thanks for flagging. I've hit this with other models before, but didn't think about it here. It would be really nice if Torch had some way to let the init function itself declare different options for different purposes

@GMNGeoffrey GMNGeoffrey May 19, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added non-zero initialization, although I actually found that the outputs weren't zero even without that.

@christinaflo christinaflo May 19, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

they wouldnt be zero bc they have the initial embeddings, but the updates should be zero if i remember correctly

@christinaflo christinaflo May 19, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@jnwei we should modify tests to use this util after it's merged in (i.e. kernel and offloading tests)

@jnwei jnwei May 21, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is a nice addition, thanks for adding this utility and for all of the testing fixes @GMNGeoffrey !

**config.architecture.heads.pairformer_embedding
).eval()

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()

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):
Expand Down