Repository: Quantum-TII/qibo (python)

Issue:
Tensorflow einsum error when measuring density matrices
**Describe the bug**
Measuring density matrices with more than 6 qubits using any Tensorflow backend leads to an einsum error. Thanks @AdrianPerezSalinas for finding this.

**To Reproduce**
```Python
from qibo import models, gates

circuit = models.Circuit(7, density_matrix=True)
circuit.add(gates.M(*range(7)))
result = circuit(nshots=100)
```
gives 

```
tensorflow.python.framework.errors_impl.UnimplementedError: Unsupported rank: 7 while handling repeated indices. Up to rank 6 is supported. [Op:Einsum]
```

**Additional context**
Einsum is used here to calculate the probability distribution from the density matrix, [by tracing out unmeasured qubits](https://github.com/Quantum-TII/qibo/blob/0422b84a189ff86dd6a900b76bd76cac90735d72/src/qibo/core/states.py#L149). Numpy backends do not have this issue.


Pull request description:
Fix einsum in density matrix traceout
Fixes #326 by modifying the einsum call used to trace-out density matrices to calculate the corresponding probability distribution for measurements.

Patch:
diff --git a/src/qibo/core/states.py b/src/qibo/core/states.py
index 63a0e6fc2..22d686cc4 100644
--- a/src/qibo/core/states.py
+++ b/src/qibo/core/states.py
@@ -54,42 +54,31 @@ class VectorState(AbstractState):
         matrix = K.outer(self.tensor, K.conj(self.tensor))
         return MatrixState.from_tensor(matrix, nqubits=self.nqubits)
 
-    def _traceout(self, qubits=None, measurement_gate=None):
-        """Helper method for :meth:`qibo.core.states.VectorState.probabilities`.
-
-        Calculates the trace-out qubit indices or einsum string.
-        """
-        if qubits is None and measurement_gate is None:
-            raise_error(ValueError, "Either ``qubits`` or ``measurement_gates`` "
-                                    "should be given to calculate measurement "
-                                    "probabilities.")
-
-        if qubits is not None:
-            if measurement_gate is not None:
+    def check_measured_qubits(func): # pylint: disable=E0213
+        """Decorator for checking list of measured qubits for probability calculation."""
+        def wrapper(self, qubits=None, measurement_gate=None):
+            if qubits is None:
+                if measurement_gate is None:
+                    raise_error(ValueError, "Either ``qubits`` or ``measurement_gates`` "
+                                            "should be given to calculate measurement "
+                                            "probabilities.")
+                if not measurement_gate.is_prepared:
+                    measurement_gate.set_nqubits(self.tensor)
+                qubits = measurement_gate.target_qubits
+            elif measurement_gate is not None:
                 raise_error(ValueError, "Cannot calculate measurement "
                                         "probabilities if both ``qubits`` and "
                                         "``measurement_gate`` are given."
                                         "Please specify only one of them.")
-            unmeasured_qubits = [i for i in range(self.nqubits)
-                                 if i not in qubits]
-            if isinstance(self, MatrixState):
-                from qibo.abstractions.callbacks import PartialTrace
-                qubits = set(unmeasured_qubits)
-                return PartialTrace.einsum_string(qubits, self.nqubits,
-                                                  measuring=True)
-            return unmeasured_qubits
-
-        if not measurement_gate.is_prepared:
-            measurement_gate.set_nqubits(self.tensor)
-        if isinstance(self, MatrixState):
-            return measurement_gate.traceout
-        return measurement_gate.unmeasured_qubits
+            return func(self, qubits=set(qubits)) # pylint: disable=E1102
+        return wrapper
 
+    @check_measured_qubits
     def probabilities(self, qubits=None, measurement_gate=None):
-        unmeasured_qubits = self._traceout(qubits, measurement_gate)
-        shape = self.nqubits * (2,)
-        state = K.reshape(K.square(K.abs(self.tensor)), shape)
-        return K.sum(state, axis=tuple(unmeasured_qubits))
+        unmeasured_qubits = tuple(i for i in range(self.nqubits)
+                                  if i not in qubits)
+        state = K.reshape(K.square(K.abs(self.tensor)), self.nqubits * (2,))
+        return K.sum(state, axis=unmeasured_qubits)
 
     def measure(self, gate, nshots, registers=None):
         self.measurements = gate(self, nshots)
@@ -144,11 +133,18 @@ class MatrixState(VectorState):
     def to_density_matrix(self):
         raise_error(RuntimeError, "State is already a density matrix.")
 
+    @VectorState.check_measured_qubits
     def probabilities(self, qubits=None, measurement_gate=None):
-        traceout = self._traceout(qubits, measurement_gate)
-        shape = 2 * self.nqubits * (2,)
-        state = K.einsum(traceout, K.reshape(self.tensor, shape))
-        return K.cast(state, dtype='DTYPE')
+        order = (tuple(sorted(qubits)) +
+                 tuple(i for i in range(self.nqubits) if i not in qubits))
+        order = order + tuple(i + self.nqubits for i in order)
+        shape = 2 * (2 ** len(qubits), 2 ** (self.nqubits - len(qubits)))
+
+        state = K.reshape(self.tensor, 2 * self.nqubits * (2,))
+        state = K.reshape(K.transpose(state, order), shape)
+        state = K.einsum("abab->a", state)
+
+        return K.reshape(K.cast(state, dtype='DTYPE'), len(qubits) * (2,))
 
 
 class DistributedState(VectorState):


Test patch:
diff --git a/src/qibo/tests_new/test_core_states.py b/src/qibo/tests_new/test_core_states.py
index 05f93453c..c78e3ea9b 100644
--- a/src/qibo/tests_new/test_core_states.py
+++ b/src/qibo/tests_new/test_core_states.py
@@ -86,36 +86,29 @@ def test_vector_state_to_density_matrix(backend):
     qibo.set_backend(original_backend)
 
 
-def test_vector_state_tracout():
+@pytest.mark.parametrize("state_type", ["VectorState", "MatrixState"])
+@pytest.mark.parametrize("use_gate", [False, True])
+def test_state_probabilities(backend, state_type, use_gate):
+    state = getattr(states, state_type).plus_state(4)
+    if use_gate:
+        from qibo import gates
+        mgate = gates.M(0, 1)
+        probs = state.probabilities(measurement_gate=mgate)
+    else:
+        probs = state.probabilities(qubits=[0, 1])
+    target_probs = np.ones((2, 2)) / 4
+    np.testing.assert_allclose(probs, target_probs)
+
+
+def test_state_probabilities_errors():
     from qibo import gates
     state = states.VectorState.zero_state(3)
     mgate = gates.M(0)
     qubits = [0]
-    assert state._traceout(qubits=qubits) == [1, 2]
-    assert state._traceout(measurement_gate=mgate) == (1, 2)
     with pytest.raises(ValueError):
-        unmeasured = state._traceout()
+        probs = state.probabilities()
     with pytest.raises(ValueError):
-        unmeasured = state._traceout(qubits, mgate)
-
-
-def test_matrix_state_tracout():
-    from qibo import gates
-    state = states.MatrixState.zero_state(2)
-    mgate = gates.M(0)
-    mgate.density_matrix = True
-    qubits = [0]
-    assert state._traceout(qubits=qubits) == "abab->a"
-    assert state._traceout(measurement_gate=mgate) == "abab->a"
-
-
-@pytest.mark.parametrize("state_type", ["VectorState", "MatrixState"])
-def test_state_probabilities(backend, state_type):
-    # TODO: Test this both for `VectorState` and `MatrixState`
-    state = getattr(states, state_type).plus_state(4)
-    probs = state.probabilities(qubits=[0, 1])
-    target_probs = np.ones((2, 2)) / 4
-    np.testing.assert_allclose(probs, target_probs)
+        probs = state.probabilities(qubits, mgate)
 
 
 @pytest.mark.parametrize("registers", [None, {"a": (0,), "b": (2,)}])