Remove unused atom14 default-frame gather

openfold/utils/feats.py

@@ -258,16 +258,15 @@ 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, ...]
 
     # [*, 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]

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