From fe219e3559edfea495b19b91cdc67cb2c2e9e1c4 Mon Sep 17 00:00:00 2001
From: Jiacheng Huang <huangjiacheng0709@outlook.com>
Date: Fri, 26 Jun 2026 16:11:17 +0800
Subject: [PATCH] test: add `softplus` operator coverage

---
 tests/test_softplus.py | 68 ++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 68 insertions(+)
 create mode 100644 tests/test_softplus.py

diff --git a/tests/test_softplus.py b/tests/test_softplus.py
new file mode 100644
index 00000000..1e63186a
--- /dev/null
+++ b/tests/test_softplus.py
@@ -0,0 +1,68 @@
+import pytest
+import torch
+import torch.nn.functional as F
+
+import infini.ops
+
+from tests.utils import Payload, empty_strided, get_stream, randn_strided
+
+
+_SHAPE_CASES = (
+    ((13, 4), None, None),
+    ((13, 4), (10, 1), (10, 1)),
+    ((13, 4, 4), None, None),
+    ((13, 4, 4), (20, 4, 1), (20, 4, 1)),
+    ((16, 5632), None, None),
+    ((4, 4, 5632), None, None),
+)
+
+_FLOAT_DTYPE_CASES = (
+    (torch.float32, 1e-6, 1e-6),
+    (torch.float16, 1e-3, 1e-3),
+    (torch.bfloat16, 1e-2, 5e-3),
+)
+
+
+@pytest.mark.auto_act_and_assert
+@pytest.mark.parametrize("shape, input_strides, out_strides", _SHAPE_CASES)
+@pytest.mark.parametrize(("dtype", "rtol", "atol"), _FLOAT_DTYPE_CASES)
+def test_softplus(
+    shape,
+    input_strides,
+    out_strides,
+    dtype,
+    device,
+    implementation_index,
+    rtol,
+    atol,
+):
+    input = randn_strided(shape, input_strides, dtype=dtype, device=device)
+    out = empty_strided(shape, out_strides, dtype=dtype, device=device)
+
+    return Payload(
+        lambda input, out: _softplus(input, out, implementation_index),
+        _torch_softplus,
+        (input, out),
+        {},
+        rtol=rtol,
+        atol=atol,
+    )
+
+
+def _softplus(input, out, implementation_index):
+    infini.ops.softplus(
+        input,
+        1.0,
+        20.0,
+        out,
+        stream=get_stream(input.device),
+        implementation_index=implementation_index,
+    )
+
+    return out
+
+
+def _torch_softplus(input, out):
+    out.copy_(F.softplus(input, beta=1.0, threshold=20.0))
+
+    return out