def _move_data(data, device):
if isinstance(data, GPUBuffer):
if device == Device.GPU: return data
old = data
data = np.empty(old.shape, dtype=np.float32)
with ProfileOp("toCPU", [data]):
cl.enqueue_copy(cl_queue, data, old.cl, is_blocking=True)
elif "ANETensor" in str(type(data)):
if device == Device.ANE: return data
with ProfileOp("toCPU", [data]):
data = data.data().astype(np.float32)
if not isinstance(data, np.ndarray):
data = np.array(data, dtype=np.float32)
if data.dtype != np.float32 and not Tensor.did_float_warning:
# warning? float64 is actually needed for numerical jacobian
print(f"warning, {data.shape!r} isn't float32")
Tensor.did_float_warning = True
if device == Device.GPU:
require_init_gpu()
with ProfileOp("toGPU", [data]):
return GPUBuffer(data.shape, data)
elif device == Device.ANE:
require_init_ane()
with ProfileOp("toANE", [data]):
ndata = ane.tensor(data.shape)
ndata.data()[:] = data
return ndata
return data
def ane(self):
assert (not self.gpu)
require_init_ane()
ndata = ane.tensor(self.shape)
ndata.data()[:] = self.data
return Tensor(ndata)