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)