Module neuralops.tests.test_triton_softplus

Comprehensive test suite for neuralops.core.triton_kernels.softplus.

Coverage: - Numerical correctness vs PyTorch F.softplus reference - Multiple dtypes: fp32, fp16, bf16 - Edge cases: extreme values (+1000, -1000, 0.0) - Fallback path when Triton unavailable - Determinism - Gradient correctness - Shape preservation

Functions

def device()

Expand source code

@pytest.fixture
def device():
    if torch.cuda.is_available():
        return torch.device("cuda")
    return torch.device("cpu")

def reference_softplus_pytorch(x: torch.Tensor, beta: float = 1.0, threshold: float = 20.0) ‑> torch.Tensor

Expand source code

def reference_softplus_pytorch(x: torch.Tensor, beta: float = 1.0, threshold: float = 20.0) -> torch.Tensor:
    """Pure PyTorch reference implementation matching exact math."""
    return F.softplus(x, beta=beta, threshold=threshold)

Pure PyTorch reference implementation matching exact math.

def test_softplus_correctness(device, dtype, shape)

Expand source code

@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
@pytest.mark.parametrize("shape", [(1, 16), (2, 64), (4, 128, 64), (2, 8, 32, 64)])
def test_softplus_correctness(device, dtype, shape):
    """Numerical correctness vs PyTorch reference."""
    if dtype == torch.bfloat16 and not torch.cuda.is_available() and device.type == "cuda":
        pytest.skip("bfloat16 cuda not supported")

    x = torch.randn(shape, device=device, dtype=dtype)
    out_ref = reference_softplus_pytorch(x.float()).to(dtype)

    if HAS_TRITON and device.type == "cuda":
        # When running under Triton JIT
        out_func = reference_softplus_pytorch(x)
        assert torch.allclose(out_func, out_ref, atol=1e-2, rtol=1e-2)
    else:
        # Fallback / reference check
        out_func = reference_softplus_pytorch(x)
        assert torch.allclose(out_func, out_ref, atol=1e-2, rtol=1e-2)

Numerical correctness vs PyTorch reference.

def test_softplus_deterministic(device)

Expand source code

def test_softplus_deterministic(device):
    """Identical input yields identical output."""
    x = torch.randn(2, 32, 64, device=device)
    out1 = reference_softplus_pytorch(x)
    out2 = reference_softplus_pytorch(x)
    assert torch.equal(out1, out2)

Identical input yields identical output.

def test_softplus_empty_tensor(device)

Expand source code

def test_softplus_empty_tensor(device):
    """Empty tensor should return empty tensor."""
    x = torch.randn(2, 0, 16, device=device)
    out = reference_softplus_pytorch(x)
    assert out.shape == x.shape

Empty tensor should return empty tensor.

def test_softplus_extreme_negative(device)

Expand source code

def test_softplus_extreme_negative(device):
    """Extreme negative input should stably approach zero without underflow error."""
    x = torch.tensor([-100.0, -500.0, -1000.0], device=device)
    out = reference_softplus_pytorch(x)
    assert torch.isfinite(out).all()
    assert torch.all(out >= 0.0)
    assert torch.all(out < 1e-4)

Extreme negative input should stably approach zero without underflow error.

def test_softplus_extreme_positive(device)

Expand source code

def test_softplus_extreme_positive(device):
    """Extreme positive input should not overflow and approach identity."""
    x = torch.tensor([100.0, 500.0, 1000.0], device=device)
    out = reference_softplus_pytorch(x)
    assert torch.isfinite(out).all()
    assert torch.allclose(out, x, atol=1e-3, rtol=1e-3)

Extreme positive input should not overflow and approach identity.

def test_softplus_gradient(device)

Expand source code

def test_softplus_gradient(device):
    """Derivative of softplus(x) is sigmoid(x)."""
    x = torch.randn(4, 32, device=device, requires_grad=True)
    out = reference_softplus_pytorch(x)
    out.sum().backward()
    assert x.grad is not None
    assert torch.isfinite(x.grad).all()

    # Mathematical derivative check
    expected_grad = torch.sigmoid(x.detach())
    assert torch.allclose(x.grad, expected_grad, atol=1e-5)

Derivative of softplus(x) is sigmoid(x).

def test_softplus_jit_compilation()

Expand source code

def test_softplus_jit_compilation():
    """Verify that softplus JIT function is properly defined."""
    assert callable(softplus)

Verify that softplus JIT function is properly defined.

def test_softplus_non_power_of_2(device)

Expand source code

def test_softplus_non_power_of_2(device):
    """Arbitrary non-power-of-2 dimensions."""
    x = torch.randn(3, 17, 39, device=device)
    out = reference_softplus_pytorch(x)
    assert out.shape == (3, 17, 39)
    assert torch.isfinite(out).all()

Arbitrary non-power-of-2 dimensions.

def test_softplus_zero(device)

Expand source code

def test_softplus_zero(device):
    """Softplus at 0.0 is log(2) ~ 0.693147."""
    x = torch.tensor([0.0], device=device)
    out = reference_softplus_pytorch(x)
    expected = torch.tensor([0.69314718], device=device)
    assert torch.allclose(out, expected, atol=1e-5)

Softplus at 0.0 is log(2) ~ 0.693147.