tests

Author: lucidrains

.github/workflows/test.yaml

@@ -0,0 +1,24 @@
+name: Tests the examples in README
+on: [push, pull_request]
+
+env:
+  TYPECHECK: True
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - name: Install Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.11"
+      - name: Install dependencies
+        run: |
+          python -m pip install uv
+          python -m uv pip install --upgrade pip
+          python -m uv pip install torch --index-url https://download.pytorch.org/whl/nightly/cpu
+          python -m uv pip install -e .[test]
+      - name: Test with pytest
+        run: |
+          python -m pytest tests/

tests/test_gaf.py

@@ -0,0 +1,68 @@
+import pytest
+from copy import deepcopy
+
+import torch
+from torch import nn
+torch.set_default_dtype(torch.float64)
+
+from GAF_microbatch_pytorch import GAFWrapper, set_filter_gradients_
+
+def test_gaf():
+
+    net = nn.Sequential(
+        nn.Linear(512, 256),
+        nn.SiLU(),
+        nn.Linear(256, 128)
+    )
+
+    gaf_net = GAFWrapper(
+        deepcopy(net),
+        filter_distance_thres = 2.
+    )
+
+    x = torch.randn(8, 1024, 512)
+    y = x.clone()
+
+    x.requires_grad_()
+    y.requires_grad_()
+
+    out1 = net(x)
+    out2 = gaf_net(y)
+
+    out1.sum().backward()
+    out2.sum().backward()
+
+    grad = net[0].weight.grad
+    grad_filtered = gaf_net.net[0].weight.grad
+
+    assert torch.allclose(grad, grad_filtered, atol = 1e-6)
+
+def test_gaf():
+
+    net = nn.Sequential(
+        nn.Linear(512, 256),
+        nn.SiLU(),
+        nn.Linear(256, 128)
+    )
+
+    gaf_net = GAFWrapper(
+        deepcopy(net),
+        filter_distance_thres = 0.
+    )
+
+    x = torch.randn(8, 1024, 512)
+    y = x.clone()
+
+    x.requires_grad_()
+    y.requires_grad_()
+
+    out1 = net(x)
+    out2 = gaf_net(y)
+
+    out1.sum().backward()
+    out2.sum().backward()
+
+    grad = net[0].weight.grad
+    grad_filtered = gaf_net.net[0].weight.grad
+
+    assert not torch.allclose(grad, grad_filtered, atol = 1e-6)