switch it back to the way Alex had before and release 0.2.0

Author: lucidrains

alphafold3_pytorch/alphafold3.py

@@ -3382,7 +3382,7 @@ def forward(
         return_present_sampled_atoms: bool = False,
         num_rollout_steps: int = 20,
         rollout_show_tqdm_pbar: bool = False
-    ) -> Float['b m 3'] | List[Float['l 3']] | Float[''] | Tuple[Float[''], LossBreakdown]:
+    ) -> Float['b m 3'] | Float['l 3'] | Float[''] | Tuple[Float[''], LossBreakdown]:
 
         atom_seq_len = atom_inputs.shape[-2]
 
@@ -3627,7 +3627,7 @@ def forward(
                 sampled_atom_pos = einx.where('b m, b m c, -> b m c', atom_mask, sampled_atom_pos, 0.)
 
             if exists(missing_atom_mask) and return_present_sampled_atoms:
-                sampled_atom_pos = [one_sampled_atom_pos[~one_missing_atom_mask] for one_sampled_atom_pos, one_missing_atom_mask in zip(sampled_atom_pos, missing_atom_mask)]
+                sampled_atom_pos = sampled_atom_pos[~missing_atom_mask]
 
             return sampled_atom_pos
 

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "alphafold3-pytorch"
-version = "0.1.136"
+version = "0.2.0"
 description = "Alphafold 3 - Pytorch"
 authors = [
     { name = "Phil Wang", email = "[email protected]" }

tests/test_input.py

@@ -191,7 +191,7 @@ def test_pdbinput_input():
 
     alphafold3.eval()
 
-    sampled_atom_pos, = alphafold3(
+    sampled_atom_pos = alphafold3(
         **batched_eval_atom_input.dict(), return_loss=False, return_present_sampled_atoms=True
     )