Repository: Quantum-TII/qibo (python)

Issue:
Cannot allocate circuits with large nqubits
As discussed today, there is a crash due to the flat shape allocation in the Circuit constructor.

Pull request description:
Move state reshape in gates
Moves the state reshape required for non-custom backends inside the gates instead of the circuit. This fixes some compatibility issues between gates of different backends and makes testing easier. 

It also fixes #337 since the `tensor_shape` and `flat_shape` attributes are no longer required during circuit creation. For this point I also added a test that initializes a QFT for n=100 qubits and checks that the depth is 2n and the number of gates n^2/2 + n.

Patch:
diff --git a/src/qibo/core/circuit.py b/src/qibo/core/circuit.py
index 2c269e3d9..65226a951 100644
--- a/src/qibo/core/circuit.py
+++ b/src/qibo/core/circuit.py
@@ -30,8 +30,6 @@ class Circuit(circuit.AbstractCircuit):
         self.param_tensor_types = K.tensor_types
         self._compiled_execute = None
         self.state_cls = states.VectorState
-        self.tensor_shape = K.cast(nqubits * (2,), dtype='DTYPEINT')
-        self.flat_shape = K.cast((2 ** nqubits,), dtype='DTYPEINT')
 
     def set_nqubits(self, gate):
         if gate.is_prepared and gate.nqubits != self.nqubits:
@@ -141,8 +139,6 @@ class Circuit(circuit.AbstractCircuit):
         """Performs all circuit gates on the state vector."""
         self._final_state = None
         state = self.get_initial_state(initial_state)
-        if not K.custom_gates:
-            state = K.reshape(state, self.tensor_shape)
 
         if self._compiled_execute is None:
             state = self._eager_execute(state)
@@ -151,9 +147,6 @@ class Circuit(circuit.AbstractCircuit):
             for callback, results in callback_results.items():
                 callback.extend(results)
 
-        if not K.custom_gates:
-            state = K.reshape(state, self.flat_shape)
-
         self._final_state = self.state_cls.from_tensor(state, self.nqubits)
         return self._final_state
 
@@ -276,8 +269,6 @@ class DensityMatrixCircuit(Circuit):
         super(DensityMatrixCircuit, self).__init__(nqubits)
         self.density_matrix = True
         self.state_cls = states.MatrixState
-        self.tensor_shape = K.cast(2 * nqubits * (2,), dtype='DTYPEINT')
-        self.flat_shape = K.cast(2 * (2 ** nqubits,), dtype='DTYPEINT')
 
     def get_initial_state(self, state=None):
         # Allow using state vectors as initial states but transform them
diff --git a/src/qibo/core/gates.py b/src/qibo/core/gates.py
index c3056a9ad..a38c661e6 100644
--- a/src/qibo/core/gates.py
+++ b/src/qibo/core/gates.py
@@ -25,6 +25,8 @@ class BackendGate(BaseBackendGate):
         from qibo.core import einsum
         self.einsum_module = einsum
         self.einsum = getattr(einsum, K.custom_einsum)()
+        self.tensor_shape = None
+        self.flat_shape = None
 
     @staticmethod
     def control_unitary(unitary):
@@ -38,32 +40,34 @@ class BackendGate(BaseBackendGate):
     def prepare(self):
         self.is_prepared = True
         self.reprepare()
-        if self.is_controlled_by:
-            if self.density_matrix:
-                # fall back to the 'defaulteinsum' backend when using
-                # density matrices with `controlled_by` gates because
-                # 'matmuleinsum' is not properly implemented for this case
-                self.einsum = self.einsum_module.DefaultEinsum()
-            self.control_cache = self.einsum_module.ControlCache(self)
-            nactive = self.nqubits - len(self.control_qubits)
-            targets = self.control_cache.targets
-            self.calculation_cache = self.einsum.create_cache(
-                targets, nactive, ncontrol=len(self.control_qubits))
-        else:
-            self.calculation_cache = self.einsum.create_cache(
-                self.qubits, self.nqubits)
-        self.calculation_cache.cast_shapes(
-            lambda x: K.cast(x, dtype='DTYPEINT'))
+        try:
+            s = 1 + self.density_matrix
+            self.tensor_shape = K.cast(s * self.nqubits * (2,), dtype='DTYPEINT')
+            self.flat_shape = K.cast(s * (2 ** self.nqubits,), dtype='DTYPEINT')
+            if self.is_controlled_by:
+                if self.density_matrix:
+                    # fall back to the 'defaulteinsum' backend when using
+                    # density matrices with `controlled_by` gates because
+                    # 'matmuleinsum' is not properly implemented for this case
+                    self.einsum = self.einsum_module.DefaultEinsum()
+                self.control_cache = self.einsum_module.ControlCache(self)
+                nactive = self.nqubits - len(self.control_qubits)
+                targets = self.control_cache.targets
+                self.calculation_cache = self.einsum.create_cache(
+                    targets, nactive, ncontrol=len(self.control_qubits))
+            else:
+                self.calculation_cache = self.einsum.create_cache(
+                    self.qubits, self.nqubits)
+            self.calculation_cache.cast_shapes(
+                lambda x: K.cast(x, dtype='DTYPEINT'))
+        except (ValueError, OverflowError):
+            pass
 
     def set_nqubits(self, state):
-        shape = tuple(state.shape)
-        if len(shape) == 1 + self.density_matrix: # pragma: no cover
-            self.nqubits = int(math.log2(shape[0]))
-        else:
-            self.nqubits = len(tuple(state.shape)) // (1 + self.density_matrix)
-        self.prepare()
+        cgates.BackendGate.set_nqubits(self, state)
 
     def state_vector_call(self, state):
+        state = K.reshape(state, self.tensor_shape)
         if self.is_controlled_by:
             ncontrol = len(self.control_qubits)
             nactive = self.nqubits - ncontrol
@@ -82,9 +86,10 @@ class BackendGate(BaseBackendGate):
         else:
             einsum_str = self.calculation_cache.vector
             state = self.einsum(einsum_str, state, self.matrix)
-        return state
+        return K.reshape(state, self.flat_shape)
 
     def density_matrix_call(self, state):
+        state = K.reshape(state, self.tensor_shape)
         if self.is_controlled_by:
             ncontrol = len(self.control_qubits)
             nactive = self.nqubits - ncontrol
@@ -111,17 +116,7 @@ class BackendGate(BaseBackendGate):
             state = self.einsum(self.calculation_cache.right, state,
                                 K.conj(self.matrix))
             state = self.einsum(self.calculation_cache.left, state, self.matrix)
-        return state
-
-    def __call__(self, state):
-        if not self.is_prepared:
-            self.set_nqubits(state)
-        original_shape = state.shape
-        if len(tuple(original_shape)) == 1 + self.density_matrix: # pragma: no cover
-            tensor_shape = (1 + self.density_matrix) * self.nqubits * (2,)
-            state = K.reshape(state, tensor_shape)
-        state = getattr(self, self._active_call)(state)
-        return K.reshape(state, original_shape)
+        return K.reshape(state, self.flat_shape)
 
 
 class H(BackendGate, gates.H):
@@ -189,18 +184,24 @@ class Collapse(BackendGate, gates.Collapse):
 
     def prepare(self):
         self.is_prepared = True
-        self.order = list(self.sorted_qubits)
-        if self.density_matrix:
-            self.order.extend((q + self.nqubits for q in self.sorted_qubits))
-            self.order.extend((q for q in range(self.nqubits)
-                               if q not in self.sorted_qubits))
-            self.order.extend((q + self.nqubits for q in range(self.nqubits)
-                               if q not in self.sorted_qubits))
-            self.sorted_qubits += [q + self.nqubits for q in self.sorted_qubits]
-            self.density_matrix_result = 2 * self.result
-        else:
-            self.order.extend((q for q in range(self.nqubits)
-                               if q not in self.sorted_qubits))
+        try:
+            self.order = list(self.sorted_qubits)
+            s = 1 + self.density_matrix
+            self.tensor_shape = K.cast(s * self.nqubits * (2,), dtype='DTYPEINT')
+            self.flat_shape = K.cast(s * (2 ** self.nqubits,), dtype='DTYPEINT')
+            if self.density_matrix:
+                self.order.extend((q + self.nqubits for q in self.sorted_qubits))
+                self.order.extend((q for q in range(self.nqubits)
+                                   if q not in self.sorted_qubits))
+                self.order.extend((q + self.nqubits for q in range(self.nqubits)
+                                   if q not in self.sorted_qubits))
+                self.sorted_qubits += [q + self.nqubits for q in self.sorted_qubits]
+                self.density_matrix_result = 2 * self.result
+            else:
+                self.order.extend((q for q in range(self.nqubits)
+                                   if q not in self.sorted_qubits))
+        except (ValueError, OverflowError): # pragma: no cover
+            pass
 
     @staticmethod
     def _append_zeros(state, qubits: List[int], results: List[int]):
@@ -216,19 +217,23 @@ class Collapse(BackendGate, gates.Collapse):
         cgates.Collapse.construct_unitary(self)
 
     def state_vector_call(self, state):
+        state = K.reshape(state, self.tensor_shape)
         substate = K.gather_nd(K.transpose(state, self.order), self.result)
         norm = K.sum(K.square(K.abs(substate)))
         state = substate / K.cast(K.sqrt(norm), dtype=state.dtype)
-        return self._append_zeros(state, self.sorted_qubits, self.result)
+        state = self._append_zeros(state, self.sorted_qubits, self.result)
+        return K.reshape(state, self.flat_shape)
 
     def density_matrix_call(self, state):
+        state = K.reshape(state, self.tensor_shape)
         substate = K.gather_nd(K.transpose(state, self.order),
                                self.density_matrix_result)
         n = 2 ** (len(tuple(substate.shape)) // 2)
         norm = K.trace(K.reshape(substate, (n, n)))
         state = substate / norm
-        return self._append_zeros(state, self.sorted_qubits,
-                                  self.density_matrix_result)
+        state = self._append_zeros(state, self.sorted_qubits,
+                                   self.density_matrix_result)
+        return K.reshape(state, self.flat_shape)
 
 
 class RX(BackendGate, gates.RX):
@@ -641,12 +646,17 @@ class _ThermalRelaxationChannelB(BackendGate, gates._ThermalRelaxationChannelB):
 
     def prepare(self):
         self.is_prepared = True
-        self.reprepare()
-        qubits = self.qubits + tuple(q + self.nqubits for q in self.qubits)
-        self.calculation_cache = self.einsum.create_cache(
-            qubits, 2 * self.nqubits)
-        self.calculation_cache.cast_shapes(
-            lambda x: K.cast(x, dtype='DTYPEINT'))
+        try:
+            self.tensor_shape = K.cast(2 * self.nqubits * (2,), dtype='DTYPEINT')
+            self.flat_shape = K.cast(2 * (2 ** self.nqubits,), dtype='DTYPEINT')
+            self.reprepare()
+            qubits = self.qubits + tuple(q + self.nqubits for q in self.qubits)
+            self.calculation_cache = self.einsum.create_cache(
+                qubits, 2 * self.nqubits)
+            self.calculation_cache.cast_shapes(
+                lambda x: K.cast(x, dtype='DTYPEINT'))
+        except (ValueError, OverflowError): # pragma: no cover
+            pass
 
     def construct_unitary(self):
         return cgates._ThermalRelaxationChannelB.construct_unitary(self)
@@ -657,4 +667,6 @@ class _ThermalRelaxationChannelB(BackendGate, gates._ThermalRelaxationChannelB):
 
     def density_matrix_call(self, state):
         einsum_str = self.calculation_cache.vector
-        return self.einsum(einsum_str, state, self.matrix)
+        state = K.reshape(state, self.tensor_shape)
+        state = self.einsum(einsum_str, state, self.matrix)
+        return K.reshape(state, self.flat_shape)


Test patch:
diff --git a/src/qibo/tests/test_density_matrix.py b/src/qibo/tests/test_density_matrix.py
index b2dbf7bd3..74147e8b9 100644
--- a/src/qibo/tests/test_density_matrix.py
+++ b/src/qibo/tests/test_density_matrix.py
@@ -20,7 +20,7 @@ def test_circuit_dm(backend):
     c.add(gates.H(1))
     c.add(gates.CNOT(0, 1))
     c.add(gates.H(2))
-    final_rho = c(np.copy(initial_rho)).numpy().reshape(initial_rho.shape)
+    final_rho = c(np.copy(initial_rho))
 
     h = np.array([[1, 1], [1, -1]]) / np.sqrt(2)
     cnot = np.array([[1, 0, 0, 0], [0, 1, 0, 0],
@@ -131,13 +131,9 @@ def test_general_channel(backend, tfmatrices, oncircuit):
     if oncircuit:
         c = models.Circuit(2, density_matrix=True)
         c.add(gate)
-        final_rho = c(np.copy(initial_rho)).numpy()
+        final_rho = c(np.copy(initial_rho))
     else:
-        if backend == "custom":
-            final_rho = gate(np.copy(initial_rho))
-        else:
-            final_rho = gate(np.copy(initial_rho).reshape(4 * (2,)))
-            final_rho = final_rho.numpy().reshape((4, 4))
+        final_rho = gate(np.copy(initial_rho))
 
     m1 = np.kron(np.eye(2), np.array(a1))
     m2 = np.array(a2)
@@ -498,12 +494,7 @@ def test_variational_layer(backend, nqubits):
     gate = gates.VariationalLayer(range(nqubits), pairs,
                                   gates.RY, gates.CZ, theta)
     gate.density_matrix = True
-    initial_state = c.get_initial_state()
-    if backend != "custom":
-        initial_state = np.reshape(initial_state, 2 * nqubits * (2,))
-    final_state = gate(initial_state)
-    if backend != "custom":
-        final_state = np.reshape(final_state, 2 * (2 ** nqubits,))
+    final_state = gate(c.get_initial_state())
     np.testing.assert_allclose(target_state, final_state)
     qibo.set_backend(original_backend)
 
diff --git a/src/qibo/tests/test_qft.py b/src/qibo/tests/test_qft.py
index 4f9f4f518..be3409b89 100644
--- a/src/qibo/tests/test_qft.py
+++ b/src/qibo/tests/test_qft.py
@@ -3,6 +3,7 @@ Testing Quantum Fourier Transform (QFT) circuit.
 """
 import numpy as np
 import pytest
+import qibo
 from qibo import gates, models
 from qibo.tests import utils
 
@@ -30,11 +31,15 @@ def exact_qft(x: np.ndarray, inverse: bool = False) -> np.ndarray:
     return qft_matrix(dim, inverse).dot(x) / np.sqrt(dim)
 
 
-def test_qft_sanity():
-    c = models.QFT(4)
-    assert c.nqubits == 4
-    assert c.depth == 8
-    assert c.ngates == 12
+@pytest.mark.parametrize("nqubits", [4, 10, 100])
+def test_qft_circuit_size(backend, nqubits):
+    original_backend = qibo.get_backend()
+    qibo.set_backend(backend)
+    c = models.QFT(nqubits)
+    assert c.nqubits == nqubits
+    assert c.depth == 2 * nqubits
+    assert c.ngates == nqubits ** 2 // 2 + nqubits
+    qibo.set_backend(original_backend)
 
 
 @pytest.mark.parametrize("nqubits", [4, 5])
diff --git a/src/qibo/tests_new/test_core_circuit_features.py b/src/qibo/tests_new/test_core_circuit_features.py
index 8a5cf3300..37e6af019 100644
--- a/src/qibo/tests_new/test_core_circuit_features.py
+++ b/src/qibo/tests_new/test_core_circuit_features.py
@@ -6,17 +6,18 @@ from qibo import gates
 from qibo.models import Circuit
 
 
-def test_circuit_vs_gate_execution(backend):
+@pytest.mark.parametrize("compile", [False, True])
+def test_circuit_vs_gate_execution(backend, compile):
     """Check consistency between executing circuit and stand alone gates."""
     from qibo import K
     original_backend = qibo.get_backend()
     qibo.set_backend(backend)
     theta = 0.1234
-    c = Circuit(2)
-    c.add(gates.X(0))
-    c.add(gates.X(1))
-    c.add(gates.CU1(0, 1, theta))
-    r1 = c.execute()
+    target_c = Circuit(2)
+    target_c.add(gates.X(0))
+    target_c.add(gates.X(1))
+    target_c.add(gates.CU1(0, 1, theta))
+    target_result = target_c()
 
     # custom circuit
     def custom_circuit(initial_state, theta):
@@ -25,21 +26,18 @@ def test_circuit_vs_gate_execution(backend):
         o = gates.CU1(0, 1, theta)(l2)
         return o
 
-    init2 = c.get_initial_state()
-    init3 = c.get_initial_state()
-    if backend != "custom":
-        init2 = K.reshape(init2, (2, 2))
-        init3 = K.reshape(init3, (2, 2))
+    initial_state = target_c.get_initial_state()
+    if compile:
+        c = K.compile(custom_circuit)
+    else:
+        c = custom_circuit
 
-    r2 = K.reshape(custom_circuit(init2, theta), (4,))
-    np.testing.assert_allclose(r1, r2)
-    compiled_custom_circuit = K.compile(custom_circuit)
-    if backend == "custom":
+    if backend == "custom" and compile:
         with pytest.raises(NotImplementedError):
-            r3 = compiled_custom_circuit(init3, theta)
+            result = c(initial_state, theta)
     else:
-        r3 = K.reshape(compiled_custom_circuit(init3, theta), (4,))
-        np.testing.assert_allclose(r2, r3)
+        result = c(initial_state, theta)
+        np.testing.assert_allclose(result, target_result)
     qibo.set_backend(original_backend)
 
 
diff --git a/src/qibo/tests_new/test_core_gates.py b/src/qibo/tests_new/test_core_gates.py
index 641b5f933..419ba5bf6 100644
--- a/src/qibo/tests_new/test_core_gates.py
+++ b/src/qibo/tests_new/test_core_gates.py
@@ -28,12 +28,8 @@ def apply_gates(gatelist, nqubits=None, initial_state=None):
                                "".format(type(initial_state)))
 
     state = K.cast(state)
-    if qibo.get_backend() != "custom":
-        state = K.qnp.reshape(state, nqubits * (2,))
     for gate in gatelist:
         state = gate(state)
-    if qibo.get_backend() != "custom":
-        state = np.array(state).ravel()
     return state
 
 
@@ -526,4 +522,5 @@ def test_callback_gate_errors():
     with pytest.raises(ValueError):
         gate.construct_unitary()
 
+
 # TODO: Test channels
diff --git a/src/qibo/tests_new/test_core_gates_density_matrix.py b/src/qibo/tests_new/test_core_gates_density_matrix.py
index a7377d19f..88def0ced 100644
--- a/src/qibo/tests_new/test_core_gates_density_matrix.py
+++ b/src/qibo/tests_new/test_core_gates_density_matrix.py
@@ -37,15 +37,10 @@ def test_hgate_density_matrix(backend):
     initial_rho = random_density_matrix(2)
     gate = gates.H(1)
     gate.density_matrix = True
-    if backend == "custom":
-        final_rho = np.copy(initial_rho)
-    else:
-        final_rho = np.copy(initial_rho).reshape(4 * (2,))
-    final_rho = np.reshape(gate(final_rho), (4, 4))
+    final_rho = gate(np.copy(initial_rho))
 
     matrix = np.array([[1, 1], [1, -1]]) / np.sqrt(2)
     matrix = np.kron(np.eye(2), matrix)
-    initial_rho = initial_rho.reshape((4, 4))
     target_rho = matrix.dot(initial_rho).dot(matrix)
     np.testing.assert_allclose(final_rho, target_rho)
     qibo.set_backend(original_backend)
@@ -156,13 +151,9 @@ def test_cu1gate_application_twoqubit(backend):
     theta = 0.1234
     nqubits = 3
     initial_rho = random_density_matrix(nqubits)
-    if backend == "custom":
-        final_rho = np.copy(initial_rho)
-    else:
-        final_rho = np.copy(initial_rho).reshape(2 * nqubits * (2,))
     gate = gates.CU1(0, 1, theta=theta)
     gate.density_matrix = True
-    final_rho = np.reshape(gate(final_rho), initial_rho.shape)
+    final_rho = gate(np.copy(initial_rho))
 
     matrix = np.eye(4, dtype=np.complex128)
     matrix[3, 3] = np.exp(1j * theta)