def test_dlpack_converter(self):
data = np.array([0., 1., 2.], 'float32')
with dragon.device('cpu'), dragon.eager_scope():
x = dragon.EagerTensor(data, copy=True)
x_to_dlpack = dragon.dlpack.to_dlpack(x)
x_from_dlpack = dragon.dlpack.from_dlpack(x_to_dlpack)
self.assertEqual(x_from_dlpack.shape, data.shape)
self.assertEqual(x_from_dlpack.dtype, str(data.dtype))
self.assertLessEqual(np.abs(x_from_dlpack.numpy() - data).max(), 1e-5)
def test_dlpack_converter_cuda(self):
data = np.array([0., 1., 2.], 'float32')
with dragon.device('cuda', 0), execution_context().mode('EAGER_MODE'):
x = dragon.EagerTensor(data, copy=True) + 0
x_to_dlpack = dragon.dlpack.to_dlpack(x)
x_from_dlpack = dragon.dlpack.from_dlpack(x_to_dlpack)
self.assertEqual(x_from_dlpack.device.type, 'cuda')
self.assertEqual(x_from_dlpack.device.index, 0)
self.assertEqual(x_from_dlpack.shape, data.shape)
self.assertEqual(x_from_dlpack.dtype, str(data.dtype))
self.assertLessEqual(np.abs(x_from_dlpack.numpy() - data).max(), 1e-5)
def test_dlpack_converter_cuda(self):
data = np.array([0., 1., 2.], 'float32')
with dragon.device('cuda', 0):
x = dragon.constant(data, copy=True) + 0
x_to_dlpack = dragon.dlpack.to_dlpack(x)
x_from_dlpack = dragon.dlpack.from_dlpack(x_to_dlpack)
self.assertEqual(x_from_dlpack.device.type, 'cuda')
self.assertEqual(x_from_dlpack.device.index, 0)
self.assertEqual(x_from_dlpack.shape, data.shape)
self.assertEqual(x_from_dlpack.dtype, str(data.dtype))
self.assertLessEqual(np.abs(x_from_dlpack.numpy() - data).max(), 1e-5)
def test_device(self):
try:
with dragon.device('abc'):
pass
except ValueError:
pass