diff --git a/tests/test_maximum.py b/tests/test_maximum.py new file mode 100644 index 00000000..9ed803fd --- /dev/null +++ b/tests/test_maximum.py @@ -0,0 +1,70 @@ +import pytest +import torch + +import infini.ops + +from tests.utils import Payload, empty_strided, get_stream, randn_strided + + +_SHAPE_CASES = ( + ((13, 4), None, None, None), + ((13, 4), (10, 1), (10, 1), (10, 1)), + ((13, 4, 4), None, None, None), + ((13, 4, 4), (20, 4, 1), (20, 4, 1), (20, 4, 1)), + ((16, 5632), None, None, None), + ((4, 4, 5632), None, 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, other_strides, out_strides", _SHAPE_CASES +) +@pytest.mark.parametrize(("dtype", "rtol", "atol"), _FLOAT_DTYPE_CASES) +def test_maximum( + shape, + input_strides, + other_strides, + out_strides, + dtype, + device, + implementation_index, + rtol, + atol, +): + input = randn_strided(shape, input_strides, dtype=dtype, device=device) + other = randn_strided(shape, other_strides, dtype=dtype, device=device) + out = empty_strided(shape, out_strides, dtype=dtype, device=device) + + return Payload( + lambda input, other, out: _maximum(input, other, out, implementation_index), + _torch_maximum, + (input, other, out), + {}, + rtol=rtol, + atol=atol, + ) + + +def _maximum(input, other, out, implementation_index): + infini.ops.maximum( + input, + other, + out, + stream=get_stream(input.device), + implementation_index=implementation_index, + ) + + return out + + +def _torch_maximum(input, other, out): + out.copy_(torch.maximum(input, other)) + + return out