test: triangular attention snapshots

openfold3/tests/test_triangular_attention.py

@@ -12,38 +12,64 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import unittest
-
+import pytest
 import torch
 
 from openfold3.core.model.layers.triangular_attention import TriangleAttention
 from openfold3.tests.config import consts
 
 
-class TestTriangularAttention(unittest.TestCase):
-    def test_shape(self):
-        c_z = consts.c_z
-        c = 12
-        no_heads = 4
-        starting = True
+# starting=True -> "starting node" variant: rows attend to rows,
+# biased by z[i, k]. False would transpose internally for the
+# "ending node" variant (columns attend to columns).
+@pytest.mark.parametrize("starting", [True, False])
+def test_shape(starting, ndarrays_regression):
+    # NOTE: seeding may need further work — torch.manual_seed controls both
+    # the random input and the module's weight init. If init changes upstream,
+    # regenerate snapshots with: pytest --force-regen
+    torch.manual_seed(123)
+
+    # c_z: pair representation channel dim (128 in production)
+    c_z = consts.c_z
+    # c: attention hidden dim (production uses 32; smaller here for speed)
+    c = 12
+    no_heads = 4
 
-        tan = TriangleAttention(
-            c_z,
-            c,
-            no_heads,
-            starting=starting,
-        )
+    tan = TriangleAttention(
+        c_z,
+        c,
+        no_heads,
+        starting=starting,
+    )
+    # AlphaFold initializes the output projection to zero (so residual blocks
+    # start as identity). Reinitialize all params so the test exercises the
+    # actual computation and produces non-trivial output.
+    for p in tan.parameters():
+        torch.nn.init.normal_(p, std=0.01)
+    tan.eval()
 
-        batch_size = consts.batch_size
-        n_res = consts.n_res
+    batch_size = consts.batch_size
+    n_res = consts.n_res
 
-        x = torch.rand((batch_size, n_res, n_res, c_z))
-        shape_before = x.shape
+    # Pair representation: [batch, N_residues, N_residues, C_z]
+    x = torch.rand((batch_size, n_res, n_res, c_z))
+    shape_before = x.shape
+    # chunk_size=None -> no memory-saving chunking, full attention in one pass
+    with torch.no_grad():
         x = tan(x, chunk_size=None)
-        shape_after = x.shape
+    shape_after = x.shape
 
-        self.assertTrue(shape_before == shape_after)
+    # Shape must be preserved for the residual addition z = z + tri_att(z)
+    assert shape_before == shape_after
 
+    # Guard against trivial all-zero output (e.g. from zero-initialized weights)
+    assert x.abs().max().item() > 0, (
+        "Output is all zeros — snapshot would be meaningless"
+    )
 
-if __name__ == "__main__":
-    unittest.main()
+    # Snapshot regression: output must be numerically identical across runs.
+    # Regenerate with: pytest --force-regen
+    ndarrays_regression.check(
+        {"output": x.cpu().numpy()},
+        default_tolerance=dict(atol=1e-6, rtol=1e-5),
+    )

openfold3/tests/test_triangular_multiplicative_update.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 
 import re
-import unittest
 
 import torch
 
@@ -23,34 +22,59 @@
 )
 from openfold3.tests.config import consts
 
+# Updates pair representation z[i,j] by projecting to two gated vectors (a, b),
+# contracting along a shared dimension (outgoing vs incoming), then projecting
+# back. "Outgoing" contracts over the starting node, "Incoming" over the ending
+# node. Shape-preserving: [*, N, N, C_z] -> [*, N, N, C_z].
 
-class TestTriangularMultiplicativeUpdate(unittest.TestCase):
-    def test_shape(self):
-        c_z = consts.c_z
-        c = 11
-
-        if re.fullmatch("^model_[1-5]_multimer_v3$", consts.model_preset):
-            tm = FusedTriangleMultiplicationOutgoing(
-                c_z,
-                c,
-            )
-        else:
-            tm = TriangleMultiplicationOutgoing(
-                c_z,
-                c,
-            )
-
-        n_res = consts.c_z
-        batch_size = consts.batch_size
-
-        x = torch.rand((batch_size, n_res, n_res, c_z))
-        mask = torch.randint(0, 2, size=(batch_size, n_res, n_res))
-        shape_before = x.shape
+
+def _make_module(c_z, c):
+    """Pick fused vs non-fused variant based on model preset."""
+    # Multimer v3 uses a fused variant (single projection split into a, b)
+    # vs separate projections for each
+    if re.fullmatch("^model_[1-5]_multimer_v3$", consts.model_preset):
+        return FusedTriangleMultiplicationOutgoing(c_z, c)
+    return TriangleMultiplicationOutgoing(c_z, c)
+
+
+def test_shape(ndarrays_regression):
+    torch.manual_seed(123)
+
+    # c_z: pair representation channel dim (128 in production)
+    c_z = consts.c_z
+    # c: hidden projection dim (production uses ~128; smaller here for speed)
+    c = 11
+
+    tm = _make_module(c_z, c)
+    # Reinitialize all params to non-trivial values (some layers may be
+    # zero-initialized by default for residual identity at init)
+    for p in tm.parameters():
+        torch.nn.init.normal_(p, std=0.01)
+    tm.eval()
+
+    n_res = consts.n_res
+    batch_size = consts.batch_size
+
+    # Pair representation: [batch, N_residues, N_residues, C_z]
+    x = torch.rand((batch_size, n_res, n_res, c_z))
+    # Binary mask: which residue pairs are valid
+    mask = torch.randint(0, 2, size=(batch_size, n_res, n_res))
+    shape_before = x.shape
+    with torch.no_grad():
         x = tm(x, mask)
-        shape_after = x.shape
+    shape_after = x.shape
 
-        self.assertTrue(shape_before == shape_after)
+    # Shape must be preserved for the residual addition z = z + tri_mul(z)
+    assert shape_before == shape_after
 
+    # Guard against trivial all-zero output (e.g. from zero-initialized weights)
+    assert x.abs().max().item() > 0, (
+        "Output is all zeros — snapshot would be meaningless"
+    )
 
-if __name__ == "__main__":
-    unittest.main()
+    # Snapshot regression: output must be numerically identical across runs.
+    # Regenerate with: pytest --force-regen
+    ndarrays_regression.check(
+        {"output": x.cpu().numpy()},
+        default_tolerance=dict(atol=1e-6, rtol=1e-5),
+    )

pyproject.toml

@@ -71,6 +71,7 @@ test = [
   "pytest-xdist",
   "pytest-cov",
   "pytest-benchmark",
+  "pytest-regressions",
   "debugpy",
 ]