fix package name so it is actually importable

Author: lucidrains

README.md

@@ -226,3 +226,14 @@ docker run -v .:/data --gpus all -it af3
     url     = {https://api.semanticscholar.org/CorpusID:268063190}
 }
 ```
+
+```bibtex
+@article{Puny2021FrameAF,
+    title   = {Frame Averaging for Invariant and Equivariant Network Design},
+    author  = {Omri Puny and Matan Atzmon and Heli Ben-Hamu and Edward James Smith and Ishan Misra and Aditya Grover and Yaron Lipman},
+    journal = {ArXiv},
+    year    = {2021},
+    volume  = {abs/2110.03336},
+    url     = {https://api.semanticscholar.org/CorpusID:238419638}
+}
+```

alphafold3_pytorch/alphafold3.py

@@ -17,7 +17,7 @@
     Sequential,
 )
 
-from typing import Literal, Tuple, NamedTuple
+from typing import Literal, Tuple, NamedTuple, Callable
 
 from alphafold3_pytorch.typing import (
     Float,
@@ -37,6 +37,8 @@
     full_pairwise_repr_to_windowed
 )
 
+from frame_averaging_pytorch import FrameAverage
+
 from taylor_series_linear_attention import TaylorSeriesLinearAttn
 
 import einx
@@ -2106,6 +2108,7 @@ def forward(
         pairwise_trunk: Float['b n n dpt'],
         pairwise_rel_pos_feats: Float['b n n dpr'],
         molecule_atom_lens: Int['b n'],
+        frame_average_fn: Callable[[Float['b n 3']], Float['b n 3']] | None = None,
         return_denoised_pos = False,
         additional_molecule_feats: Float[f'b n {ADDITIONAL_MOLECULE_FEATS}'] | None = None,
         add_smooth_lddt_loss = False,
@@ -2142,6 +2145,13 @@ def forward(
             )
         )
 
+        # frame average the denoised atom positions if needed
+
+        if exists(frame_average_fn):
+            denoised_atom_pos = frame_average_fn(denoised_atom_pos)
+
+        # total loss, for accumulating all auxiliary losses
+
         total_loss = 0.
 
         # if additional molecule feats is provided

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "alphafold3-pytorch"
-version = "0.1.16"
+version = "0.1.17"
 description = "Alphafold 3 - Pytorch"
 authors = [
     { name = "Phil Wang", email = "[email protected]" }
@@ -29,7 +29,7 @@ dependencies = [
     "einx>=0.2.2",
     "ema-pytorch>=0.4.8",
     "environs",
-    "frame-averaging-pytorch",
+    "frame-averaging-pytorch>=0.0.17",
     "hydra-core",
     "jaxtyping>=0.2.28",
     "lightning>=2.2.5",
@@ -75,4 +75,4 @@ ignore = [
 allow-direct-references = true
 
 [tool.hatch.build.targets.wheel]
-packages = ["alphafold3-pytorch"]
+packages = ["alphafold3_pytorch"]