FlagGems/benchmark/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
import pytest
import torch

from . import base, consts, utils


def _tensor_input_fn(shape, dtype, device):
    inp1 = utils.generate_tensor_input(shape, dtype, device)
    inp2 = utils.generate_tensor_input(shape, dtype, device)
    inp2 = torch.where(inp2 == 0, torch.ones_like(inp2), inp2)
    yield inp1, inp2


def _scalar_input_fn(shape, dtype, device):
    inp1 = utils.generate_tensor_input(shape, dtype, device)
    yield inp1, 0.5


@pytest.mark.fmod
def test_fmod_tensor():
    bench = base.GenericBenchmark(
        input_fn=_tensor_input_fn,
        op_name="fmod.Tensor",
        torch_op=torch.fmod,
        dtypes=consts.FLOAT_DTYPES,
    )
    bench.run()


@pytest.mark.fmod
def test_fmod_scalar():
    bench = base.GenericBenchmark(
        input_fn=_scalar_input_fn,
        op_name="fmod.Scalar",
        torch_op=torch.fmod,
        dtypes=consts.FLOAT_DTYPES,
    )
    bench.run()


@pytest.mark.fmod
def test_fmod_tensor_():
    bench = base.GenericBenchmark(
        input_fn=_tensor_input_fn,
        op_name="fmod_.Tensor",
        torch_op=torch.Tensor.fmod_,
        dtypes=consts.FLOAT_DTYPES,
        inplace=True,
    )
    bench.run()


@pytest.mark.fmod
def test_fmod_scalar_():
    bench = base.GenericBenchmark(
        input_fn=_scalar_input_fn,
        op_name="fmod_.Scalar",
        torch_op=torch.Tensor.fmod_,
        dtypes=consts.FLOAT_DTYPES,
        inplace=True,
    )
    bench.run()