FlagGems/tests/test_fmod.py at e1dc6e6526095c139166c3269538bee227dc86a7 · Schopenhauer-loves-Hegel/FlagGems · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
import pytest
import torch

import flag_gems

from .accuracy_utils import (
    FLOAT_DTYPES,
    POINTWISE_SHAPES,
    SCALARS,
    gems_assert_close,
    to_reference,
)


@pytest.mark.fmod
@pytest.mark.parametrize("shape", POINTWISE_SHAPES)
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
def test_accuracy_fmod(shape, dtype):
    inp1 = torch.randn(shape, dtype=dtype, device=flag_gems.device)
    inp2 = torch.randn(shape, dtype=dtype, device=flag_gems.device)
    inp2 = torch.where(inp2 == 0, torch.ones_like(inp2), inp2)
    ref_inp1 = to_reference(inp1, True)
    ref_inp2 = to_reference(inp2, True)
    ref_out = torch.fmod(ref_inp1, ref_inp2)
    with flag_gems.use_gems():
        res_out = torch.fmod(inp1, inp2)
    gems_assert_close(res_out, ref_out, dtype)


@pytest.mark.fmod
@pytest.mark.parametrize("shape", POINTWISE_SHAPES)
@pytest.mark.parametrize("scalar", SCALARS)
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
def test_accuracy_fmod_tensor_scalar(shape, scalar, dtype):
    inp1 = torch.randn(shape, dtype=dtype, device=flag_gems.device)
    inp2 = scalar if scalar != 0 else 1.0
    ref_inp1 = to_reference(inp1, True)
    ref_out = torch.fmod(ref_inp1, inp2)
    with flag_gems.use_gems():
        res_out = torch.fmod(inp1, inp2)
    atol = 1e-3 if abs(scalar) < 0.01 else 1e-4
    gems_assert_close(res_out, ref_out, dtype, atol=atol)


@pytest.mark.inplace
@pytest.mark.fmod_
@pytest.mark.parametrize("shape", POINTWISE_SHAPES)
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
def test_accuracy_fmod_(shape, dtype):
    inp1 = torch.randn(shape, dtype=dtype, device=flag_gems.device)
    inp2 = torch.randn(shape, dtype=dtype, device=flag_gems.device)
    inp2 = torch.where(inp2 == 0, torch.ones_like(inp2), inp2)
    ref_inp1 = to_reference(inp1.clone(), True)
    ref_inp2 = to_reference(inp2, True)
    ref_out = ref_inp1.fmod_(ref_inp2)
    with flag_gems.use_gems():
        res_out = inp1.fmod_(inp2)
    gems_assert_close(res_out, ref_out, dtype)


@pytest.mark.inplace
@pytest.mark.fmod_
@pytest.mark.parametrize("shape", POINTWISE_SHAPES)
@pytest.mark.parametrize("scalar", SCALARS)
@pytest.mark.parametrize("dtype", FLOAT_DTYPES)
def test_fmod_scalar_inplace(shape, scalar, dtype):
    inp1 = torch.randn(shape, dtype=dtype, device=flag_gems.device)
    inp2 = scalar if scalar != 0 else 1.0
    ref_inp1 = to_reference(inp1.clone(), True)
    ref_out = ref_inp1.fmod_(inp2)
    with flag_gems.use_gems():
        res_out = inp1.fmod_(inp2)
    atol = 1e-3 if abs(scalar) < 0.01 else 1e-4
    gems_assert_close(res_out, ref_out, dtype, atol=atol)