pytorch/caffe2/python/operator_test/cast_op_test.py
Bugra Akyildiz 27c7158166 Remove __future__ imports for legacy Python2 supports (#45033)
Summary:
There is a module called `2to3` which you can target for future specifically to remove these, the directory of `caffe2` has the most redundant imports:

```2to3 -f future -w caffe2```

Pull Request resolved: https://github.com/pytorch/pytorch/pull/45033

Reviewed By: seemethere

Differential Revision: D23808648

Pulled By: bugra

fbshipit-source-id: 38971900f0fe43ab44a9168e57f2307580d36a38
2020-09-23 17:57:02 -07:00

48 lines
1.6 KiB
Python

from caffe2.python import core, workspace
import caffe2.python.hypothesis_test_util as hu
from hypothesis import given
import numpy as np
class TestCastOp(hu.HypothesisTestCase):
@given(**hu.gcs)
def test_cast_int_float(self, gc, dc):
data = np.random.rand(5, 5).astype(np.int32)
# from int to float
op = core.CreateOperator('Cast', 'data', 'data_cast', to=1, from_type=2)
self.assertDeviceChecks(dc, op, [data], [0])
# This is actually 0
self.assertGradientChecks(gc, op, [data], 0, [0])
@given(**hu.gcs)
def test_cast_int_float_empty(self, gc, dc):
data = np.random.rand(0).astype(np.int32)
# from int to float
op = core.CreateOperator('Cast', 'data', 'data_cast', to=1, from_type=2)
self.assertDeviceChecks(dc, op, [data], [0])
# This is actually 0
self.assertGradientChecks(gc, op, [data], 0, [0])
@given(data=hu.tensor(dtype=np.int32), **hu.gcs_cpu_only)
def test_cast_int_to_string(self, data, gc, dc):
op = core.CreateOperator(
'Cast', 'data', 'data_cast', to=core.DataType.STRING)
def ref(data):
ret = data.astype(dtype=np.str)
# the string blob will be fetched as object, we feed and re-fetch
# to mimic this.
with hu.temp_workspace('tmp_ref_int_to_string'):
workspace.FeedBlob('tmp_blob', ret)
fetched_ret = workspace.FetchBlob('tmp_blob')
return (fetched_ret, )
self.assertReferenceChecks(gc, op, inputs=[data], reference=ref)