able to return the inverse function from mean_pool_fixed_windows_with_mask for token -> atom feats

Author: lucidrains

alphafold3_pytorch/alphafold3.py

@@ -297,8 +297,8 @@ def mean_pool_fixed_windows_with_mask(
     feats: Float['b m d'],
     mask: Bool['b m'],
     window_size: int,
-    return_pooled_mask: bool = False
-) -> Float['b n d'] | Tuple[Float['b n'], Bool['b n']]:
+    return_mask_and_inverse: bool = False,
+) -> Float['b n d'] | Tuple[Float['b n d'], Bool['b n'], Callable[[Float['b m d']], Float['b n d']]]:
 
     seq_len = feats.shape[-2]
     assert divisible_by(seq_len, window_size)
@@ -310,11 +310,18 @@ def mean_pool_fixed_windows_with_mask(
 
     avg = num / den.clamp(min = 1.)
 
-    if not return_pooled_mask:
+    if not return_mask_and_inverse:
         return avg
 
     pooled_mask = reduce(mask, 'b (n w) -> b n', 'any', w = window_size)
-    return avg, pooled_mask
+
+    @typecheck
+    def inverse_fn(pooled: Float['b n d']) -> Float['b m d']:
+        unpooled = repeat(pooled, 'b n d -> b (n w) d', w = window_size)
+        unpooled = einx.where('b m, b m d, -> b m d', mask, unpooled, 0.)
+        return unpooled
+
+    return avg, pooled_mask, inverse_fn
 
 @typecheck
 def batch_repeat_interleave(

pyproject.toml

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

tests/test_af3.py

@@ -80,8 +80,9 @@ def test_mean_pool_with_mask():
     seq = torch.tensor([[[1.], [100.], [1.], [2.], [2.], [100.], [1.], [1.], [100.]]])
     mask = torch.tensor([[True, False, True, True, True, False, True, True, False]])
 
-    pooled = mean_pool_fixed_windows_with_mask(seq, mask, window_size = 3)
+    pooled, _, inverse_function = mean_pool_fixed_windows_with_mask(seq, mask, window_size = 3, return_mask_and_inverse = True)
 
+    assert inverse_function(pooled).shape == seq.shape
     assert torch.allclose(pooled, torch.tensor([[[1.], [2.], [1.]]]))
 
 def test_batch_repeat_interleave():