pytorch

1
#!/usr/bin/env python3
2

3
import caffe2.python.hypothesis_test_util as hu
4
import hypothesis.strategies as st
5
import numpy as np
6
from caffe2.python import core
7
from hypothesis import given
8

9

10
class TestAsyncNetBarrierOp(hu.HypothesisTestCase):
11
    @given(
12
        n=st.integers(1, 5),
13
        shape=st.lists(st.integers(0, 5), min_size=1, max_size=3),
14
        **hu.gcs
15
    )
16
    def test_async_net_barrier_op(self, n, shape, dc, gc):
17
        test_inputs = [(100 * np.random.random(shape)).astype(np.float32) for _ in range(n)]
18
        test_input_blobs = ["x_{}".format(i) for i in range(n)]
19

20
        barrier_op = core.CreateOperator(
21
            "AsyncNetBarrier",
22
            test_input_blobs,
23
            test_input_blobs,
24
            device_option=gc,
25
        )
26

27
        def reference_func(*args):
28
            self.assertEqual(len(args), n)
29
            return args
30

31
        self.assertReferenceChecks(gc, barrier_op, test_inputs, reference_func)
32