pytorch

1

2

3

4

5

6
import hypothesis.strategies as st
7
import numpy as np
8

9
from caffe2.python import core
10
from hypothesis import given, settings
11
import caffe2.python.hypothesis_test_util as hu
12
import caffe2.python.serialized_test.serialized_test_util as serial
13

14

15
class TestArgOps(serial.SerializedTestCase):
16
    @given(
17
        X=hu.tensor(dtype=np.float32), axis=st.integers(-1, 5),
18
        keepdims=st.booleans(), **hu.gcs)
19
    @settings(deadline=None)
20
    def test_argmax(self, X, axis, keepdims, gc, dc):
21
        if axis >= len(X.shape):
22
            axis %= len(X.shape)
23
        op = core.CreateOperator(
24
            "ArgMax", ["X"], ["Indices"], axis=axis, keepdims=keepdims,
25
            device_option=gc)
26

27
        def argmax_ref(X):
28
            indices = np.argmax(X, axis=axis)
29
            if keepdims:
30
                out_dims = list(X.shape)
31
                out_dims[axis] = 1
32
                indices = indices.reshape(tuple(out_dims))
33
            return [indices]
34

35
        self.assertReferenceChecks(gc, op, [X], argmax_ref)
36
        self.assertDeviceChecks(dc, op, [X], [0])
37

38
    @given(
39
        X=hu.tensor(dtype=np.float32), axis=st.integers(-1, 5),
40
        keepdims=st.booleans(), **hu.gcs)
41
    @settings(deadline=None)
42
    def test_argmin(self, X, axis, keepdims, gc, dc):
43
        if axis >= len(X.shape):
44
            axis %= len(X.shape)
45
        op = core.CreateOperator(
46
            "ArgMin", ["X"], ["Indices"], axis=axis, keepdims=keepdims,
47
            device_option=gc)
48

49
        def argmin_ref(X):
50
            indices = np.argmin(X, axis=axis)
51
            if keepdims:
52
                out_dims = list(X.shape)
53
                out_dims[axis] = 1
54
                indices = indices.reshape(tuple(out_dims))
55
            return [indices]
56

57
        self.assertReferenceChecks(gc, op, [X], argmin_ref)
58
        self.assertDeviceChecks(dc, op, [X], [0])
59

60

61
if __name__ == "__main__":
62
    import unittest
63
    unittest.main()
64