From 29f7dde328442b7900ad5cbb87fc0a0f04701183 Mon Sep 17 00:00:00 2001
From: taivu1998 <46636857+taivu1998@users.noreply.github.com>
Date: Sun, 10 May 2026 04:39:50 -0700
Subject: [PATCH] Remove unused atom14 default frame gather

---
 openfold/utils/feats.py |  5 ++---
 tests/test_feats.py     | 33 +++++++++++++++++++++++++++++++++
 2 files changed, 35 insertions(+), 3 deletions(-)

diff --git a/openfold/utils/feats.py b/openfold/utils/feats.py
index 75845d0fe..2213ea0a4 100644
--- a/openfold/utils/feats.py
+++ b/openfold/utils/feats.py
@@ -258,8 +258,7 @@ def frames_and_literature_positions_to_atom14_pos(
     atom_mask,
     lit_positions,
 ):
-    # [*, N, 14, 4, 4]
-    default_4x4 = default_frames[aatype, ...]
+    num_rigid_groups = default_frames.shape[-3]
 
     # [*, N, 14]
     group_mask = group_idx[aatype, ...]
@@ -267,7 +266,7 @@ def frames_and_literature_positions_to_atom14_pos(
     # [*, N, 14, 8]
     group_mask = nn.functional.one_hot(
         group_mask,
-        num_classes=default_frames.shape[-3],
+        num_classes=num_rigid_groups,
     )
 
     # [*, N, 14, 8]
diff --git a/tests/test_feats.py b/tests/test_feats.py
index 7a1783b85..d28fa43e3 100644
--- a/tests/test_feats.py
+++ b/tests/test_feats.py
@@ -341,6 +341,39 @@ def test_frames_and_literature_positions_to_atom14_pos_shape(self):
 
         self.assertTrue(xyz.shape == (batch_size, n_res, 14, 3))
 
+    def test_atom14_pos_does_not_gather_default_frames(self):
+        class DefaultFrameShapeOnly:
+            shape = torch.Size((21, 8, 4, 4))
+
+            def __getitem__(self, key):
+                raise AssertionError("default_frames should not be gathered")
+
+        batch_size = 1
+        n_res = 3
+
+        rots = torch.eye(3).expand(batch_size, n_res, 8, 3, 3)
+        trans = torch.zeros((batch_size, n_res, 8, 3))
+
+        if consts.is_multimer:
+            rotation = Rot3Array.from_array(rots)
+            translation = Vec3Array.from_array(trans)
+            ts = Rigid3Array(rotation, translation)
+        else:
+            ts = Rigid(Rotation(rot_mats=rots), trans)
+
+        aatype = torch.randint(low=0, high=21, size=(batch_size, n_res)).long()
+
+        xyz = feats.frames_and_literature_positions_to_atom14_pos(
+            ts,
+            aatype,
+            DefaultFrameShapeOnly(),
+            torch.tensor(restype_atom14_to_rigid_group),
+            torch.tensor(restype_atom14_mask),
+            torch.tensor(restype_atom14_rigid_group_positions),
+        )
+
+        self.assertTrue(xyz.shape == (batch_size, n_res, 14, 3))
+
     @compare_utils.skip_unless_alphafold_installed()
     def test_frames_and_literature_positions_to_atom14_pos_compare(self):
         def run_f(aatype, affines):