From 9ecff0d361d6121fdbe98c07b523abd49c399575 Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Fri, 26 Jun 2026 16:11:32 +0800 Subject: [PATCH] test: add `threshold` operator coverage --- tests/test_threshold.py | 68 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 68 insertions(+) create mode 100644 tests/test_threshold.py diff --git a/tests/test_threshold.py b/tests/test_threshold.py new file mode 100644 index 00000000..55b977bd --- /dev/null +++ b/tests/test_threshold.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, 0.0, 0.0), + (torch.float16, 0.0, 0.0), + (torch.bfloat16, 0.0, 0.0), +) + + +@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_threshold( + 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: _threshold(input, out, implementation_index), + _torch_threshold, + (input, out), + {}, + rtol=rtol, + atol=atol, + ) + + +def _threshold(input, out, implementation_index): + infini.ops.threshold( + input, + 0.25, + -1.0, + out, + stream=get_stream(input.device), + implementation_index=implementation_index, + ) + + return out + + +def _torch_threshold(input, out): + out.copy_(F.threshold(input, threshold=0.25, value=-1.0)) + + return out