def __init__(self, circuit): from qibo.core.distcircuit import DistributedCircuit super().__init__(circuit.nqubits) self.circuit_cls = DistributedCircuit if not isinstance(circuit, self.circuit_cls): raise_error( TypeError, "Circuit of unsupported type {} was given to " "distributed state.".format(type(circuit))) self.circuit = circuit # List of length ``ndevices`` holding ``tf.Variable``s with the state pieces self.pieces = None # Dictionaries containing tensors that are useful for reshaping # the state when splitting or merging the pieces. n = self.nstates // 2**self.nglobal self.shapes = { "full": K.cpu_cast((self.nstates, ), dtype='DTYPEINT'), "device": K.cpu_cast((self.ndevices, n), dtype='DTYPEINT'), "tensor": self.nqubits * (2, ) } self.bintodec = { "global": 2**K.np.arange(self.nglobal - 1, -1, -1), "local": 2**K.np.arange(self.nlocal - 1, -1, -1) }
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])