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])
