pytest

Author: lucidrains

.github/workflows/test.yml

@@ -0,0 +1,21 @@
+name: Pytest
+on: [push, pull_request]
+
+jobs:
+  build:
+
+    runs-on: ubuntu-latest
+
+    steps:
+    - uses: actions/checkout@v4
+    - name: Set up Python 3.10
+      uses: actions/setup-python@v5
+      with:
+        python-version: "3.10"
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        python -m pip install -e .[test]
+    - name: Test with pytest
+      run: |
+        python -m pytest tests/

tests/test_sparse_attn.py

@@ -0,0 +1,21 @@
+import pytest
+import torch
+
+def test_sparse_attn():
+    from native_sparse_attention_pytorch import SparseAttention
+
+    attn = SparseAttention(
+        dim = 512,
+        dim_head = 64,
+        heads = 8,
+        sliding_window_size = 2,
+        compress_block_size = 4,
+        selection_block_size = 4,
+        num_selected_blocks = 2
+    )
+
+    tokens = torch.randn(2, 31, 512)
+
+    attended = attn(tokens)
+
+    assert tokens.shape == attended.shape