pytorch/caffe2/python/operator_test/bucketize_op_test.py
Richard Barnes 9945fd7253 Drop unused imports from caffe2/python (#49980)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/49980

From
```
./python/libcst/libcst codemod remove_unused_imports.RemoveUnusedImportsWithGlean --no-format caffe2/
```

Test Plan: Standard sandcastle tests

Reviewed By: xush6528

Differential Revision: D25727359

fbshipit-source-id: c4f60005b10546423dc093d31d46deb418352286
2021-01-05 13:17:46 -08:00

34 lines
930 B
Python

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()