Module neuralops.tests.test_block_sparse

Unit tests for block-sparse attention.

Functions

def test_block_sparse_attention_forward()

Expand source code

def test_block_sparse_attention_forward():
    q = torch.randn(2, 4, 64, 32)
    k = torch.randn(2, 4, 64, 32)
    v = torch.randn(2, 4, 64, 32)
    out = block_sparse_attention(q, k, v, block_size=16)
    assert out.shape == (2, 4, 64, 32)
    assert not torch.isnan(out).any()

def test_butterfly_mask_shape_and_diag()

Expand source code

def test_butterfly_mask_shape_and_diag():
    mask = butterfly_mask(128, block_size=32)
    assert mask.shape == (128, 128)
    # Diagonal must be True
    assert mask.diag().all()