Module neuralops.tests.test_kernel_quantization_comprehensive

Functions

def main()

Expand source code

def main():
    print("Testing Attentions...")
    config = MockConfig()
    for attn_cls, name in [
        (ClusterAttention, "ClusterAttention"),
        (AdaSplashAttention, "AdaSplashAttention"),
        (HammingAttention, "HammingAttention"),
        (GalerkinAttention, "GalerkinAttention"),
        (PerformerAttention, "PerformerAttention"),
    ]:
        try:
            _run_attention_test(attn_cls(config), name)
        except Exception as e:
            print(f"{name} failed: {e}")
            traceback.print_exc()

def test_all_attentions(attn_cls, name)

Expand source code

@pytest.mark.parametrize("attn_cls,name", [
    (ClusterAttention, "ClusterAttention"),
    (AdaSplashAttention, "AdaSplashAttention"),
    (HammingAttention, "HammingAttention"),
    (GalerkinAttention, "GalerkinAttention"),
    (PerformerAttention, "PerformerAttention"),
])
def test_all_attentions(attn_cls, name):
    config = MockConfig()
    module = attn_cls(config)
    _run_attention_test(module, name)