def assign_pieces(self, tensor):
"""Assigns state pieces from a given full state vector.
Args:
tensor (K.Tensor): The full state vector as a tensor supported by
the underlying backend.
"""
if self.pieces is None:
self.create_pieces()
with K.on_cpu():
tensor = K.reshape(K.cpu_cast(tensor), self.shapes["device"])
pieces = [tensor[i] for i in range(self.ndevices)]
new_tensor = K.zeros(self.shapes["device"])
with K.on_cpu():
new_tensor = K.transpose_state(pieces, new_tensor, self.nqubits,
self.qubits.transpose_order)
for i in range(self.ndevices):
K.cpu_assign(self, i, new_tensor[i])
def device_job(ids, device):
for i in ids:
piece = self._apply_gates(state.pieces[i], queues[i], device)
K.cpu_assign(state, i, piece)
del (piece)
