Module neuralops.tests.test_elementwise_kernels

Tests for elementwise kernels.

Functions

def test_fused_fast_gelu()

Expand source code

def test_fused_fast_gelu():
    x = torch.randn(2, 8)
    out = fused_fast_gelu(x)
    ref = F.gelu(x, approximate="tanh")
    assert torch.allclose(out, ref, atol=1e-4)

def test_fused_swiglu()

Expand source code

def test_fused_swiglu():
    x = torch.randn(2, 8)
    gate = torch.randn(2, 8)
    out = fused_swiglu(x, gate)
    assert torch.allclose(out, F.silu(gate) * x)