pytorch

1
import numpy
2

3
from caffe2.python import core
4
from hypothesis import given, settings
5

6
import caffe2.python.hypothesis_test_util as hu
7
import hypothesis.strategies as st
8
import numpy as np
9

10

11
@st.composite
12
def _data(draw):
13
    return draw(
14
        hu.tensor(
15
            dtype=np.int64,
16
            elements=st.integers(
17
                min_value=np.iinfo(np.int64).min, max_value=np.iinfo(np.int64).max
18
            )
19
        )
20
    )
21

22

23
class TestMod(hu.HypothesisTestCase):
24
    @settings(deadline=None)
25
    @given(
26
        data=_data(),
27
        divisor=st.integers(
28
            min_value=np.iinfo(np.int64).min, max_value=np.iinfo(np.int64).max
29
        ),
30
        inplace=st.booleans(),
31
        sign_follow_divisor=st.booleans(),
32
        **hu.gcs
33
    )
34
    def test_mod(
35
        self, data, divisor, inplace, sign_follow_divisor, gc, dc
36
    ):
37
        if divisor == 0:
38
            # invalid test case
39
            return None
40

41
        def ref(data):
42
            if sign_follow_divisor:
43
                output = data % divisor
44
            else:
45
                output = numpy.fmod(data, divisor)
46
            return [output]
47

48
        op = core.CreateOperator(
49
            'Mod',
50
            ['data'],
51
            ['data' if inplace else 'output'],
52
            divisor=divisor,
53
            sign_follow_divisor=sign_follow_divisor
54
        )
55

56
        self.assertReferenceChecks(gc, op, [data], ref)
57

58

59
if __name__ == "__main__":
60
    import unittest
61
    unittest.main()
62