def _move_kernel_to_cpu(kernel):
"""Moves data in a kernel from an accelerator to the CPU."""
if isinstance(kernel, Kernel):
return Kernel(device_get(kernel.var1), device_get(kernel.nngp),
device_get(kernel.var2), device_get(kernel.ntk),
kernel.is_gaussian, kernel.is_height_width)
else:
return device_get(kernel)
def _move_kernel_to_cpu(k):
"""Moves data in a kernel from an accelerator to the CPU."""
if hasattr(k, '_asdict'):
return k._replace(**dict([(k, device_get(v))
for k, v in k._asdict().items()]))
elif isinstance(k, np.ndarray):
return device_get(k)
else:
raise TypeError(
'Expected kernel to be either a namedtuple or a `np.ndarray`, got %s.'
% type(k))
def test_device_put_and_get(self):
x = onp.arange(12.).reshape((3, 4)).astype("float32")
dx = api.device_put(x)
self.assertIsInstance(dx, DeviceArray)
x2 = api.device_get(dx)
self.assertIsInstance(x2, onp.ndarray)
assert onp.all(x == x2)
y = [x, (2 * x, 3 * x)]
dy = api.device_put(y)
y2 = api.device_get(dy)
self.assertIsInstance(y2, list)
self.assertIsInstance(y2[0], onp.ndarray)
assert onp.all(y2[0] == x)
self.assertIsInstance(y2[1], tuple)
self.assertIsInstance(y2[1][0], onp.ndarray)
assert onp.all(y2[1][0] == 2 * x)
self.assertIsInstance(y2[1][1], onp.ndarray)
assert onp.all(y2[1][1] == 3 * x)
def x_cpu():
return device_get(
random.normal(random.PRNGKey(1), (2, 3), dtype))
def kernel_fn(x1, x2=None, *args, **kwargs): return device_get(_kernel_fn(x1, x2, *args, **kwargs))
def transform(c, rng=None):
return DataChunk(device_get(c.X, np.float64),
device_get(c.Y, np.float64), c.image_size,
c.image_channels, c.label_dim, c.label_format)
