token bond should probably be an embedding, for special bonds like disulfides.. fix later, but change to float for now

Author: lucidrains

alphafold3_pytorch/alphafold3.py

@@ -2804,7 +2804,7 @@ def forward(
         additional_residue_feats: Float['b n 10'],
         residue_atom_lens: Int['b n'] | None = None,
         atom_mask: Bool['b m'] | None = None,
-        token_bond: Bool['b n n'] | None = None,
+        token_bond: Float['b n n'] | None = None,
         msa: Float['b s n d'] | None = None,
         msa_mask: Bool['b s'] | None = None,
         templates: Float['b t n n dt'] | None = None,
@@ -2882,14 +2882,14 @@ def forward(
             # (1) mask out diagonal - token to itself does not count as a bond
             # (2) symmetrize, in case it is not already symmetrical (could also throw an error)
 
-            token_bond = token_bond | rearrange(token_bond, 'b i j -> b j i')
+            assert torch.allclose(token_bond, rearrange(token_bond, 'b i j -> b j i')), 'token bond must be symmetrical'
             diagonal = torch.eye(seq_len, device = self.device, dtype = torch.bool)
-            token_bond.masked_fill_(diagonal, False)
+            token_bond.masked_fill_(diagonal, 0.)
         else:
             seq_arange = torch.arange(seq_len, device = self.device)
-            token_bond = einx.subtract('i, j -> i j', seq_arange, seq_arange).abs() == 1
+            token_bond = (einx.subtract('i, j -> i j', seq_arange, seq_arange).abs() == 1).float()
 
-        token_bond_feats = self.token_bond_to_pairwise_feat(token_bond.float())
+        token_bond_feats = self.token_bond_to_pairwise_feat(token_bond)
 
         pairwise_init = pairwise_init + token_bond_feats
 

pyproject.toml

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

tests/test_af3.py

@@ -366,7 +366,8 @@ def test_alphafold3():
     seq_len = 16
     atom_seq_len = seq_len * 27
 
-    token_bond = torch.randint(0, 2, (2, seq_len, seq_len)).bool()
+    token_bond = torch.randint(0, 2, (2, seq_len, seq_len)).float()
+    token_bond = token_bond + token_bond.transpose(-1, -2)
 
     atom_inputs = torch.randn(2, atom_seq_len, 77)
     atom_mask = torch.ones((2, atom_seq_len)).bool()
@@ -512,7 +513,8 @@ def test_alphafold3_with_packed_atom_repr():
 
     atom_seq_len = residue_atom_lens.sum(dim = -1).amax()
 
-    token_bond = torch.randint(0, 2, (2, seq_len, seq_len)).bool()
+    token_bond = torch.randint(0, 2, (2, seq_len, seq_len)).float()
+    token_bond = token_bond + token_bond.transpose(-1, -2)
 
     atom_inputs = torch.randn(2, atom_seq_len, 77)