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
|
from caffe2.python import core
from hypothesis import given
import caffe2.python.hypothesis_test_util as hu
import numpy as np
class TestBucketizeOp(hu.HypothesisTestCase):
@given(
x=hu.tensor(
min_dim=1, max_dim=2, dtype=np.float32,
elements=hu.floats(min_value=-5, max_value=5)),
**hu.gcs)
def test_bucketize_op(self, x, gc, dc):
length = np.random.randint(low=1, high=5)
boundaries = np.random.randn(length) * 5
boundaries.sort()
def ref(x, boundaries):
bucket_idx = np.digitize(x, boundaries, right=True)
return [bucket_idx]
op = core.CreateOperator('Bucketize',
["X"], ["INDICES"],
boundaries=boundaries)
self.assertReferenceChecks(gc, op, [x, boundaries], ref)
if __name__ == "__main__":
import unittest
unittest.main()
|
