Repository: Quantum-TII/qibo (python)

Issue:
Controlled by gates with density matrices
**Describe the bug**
Controlled gates that are created using the `controlled_by` method do not work with density matrices of 5 or more qubits if the `defaulteinsum` or `matmuleinsum` Tensorflow backends are used.

**To Reproduce**
```Python
import qibo
from qibo import gates, models
qibo.set_backend("defaulteinsum")

c = models.Circuit(5, density_matrix=True)
c.add(gates.H(1).controlled_by(0))
result = c()
```
leads to the following Tensorflow error:

```tensorflow.python.framework.errors_impl.UnimplementedError: Unhandled input dimensions 10 [Op:StridedSlice] name: strided_slice/```

The issue appears only for 5 or more qubits and only with the native tensorflow backends only. The numpy backends and the custom backend work without problem.

Pull request description:
Fix issue with tf slicing
Fixes #342. In principle this is a tf issue as the existing code in master still works with numpy, I just changed the slicing to `K.gather` so that it works with Tensorflow too.

I think the problem of #342 should now be fixed but I would like to do a few more tests with more qubits because based on some other similar tf issues we fixed recently I suspect the tf einsum may fail in some cases. I am not 100% certain though.

Patch:
diff --git a/src/qibo/backends/abstract.py b/src/qibo/backends/abstract.py
index 7f78035ae..b0d846409 100644
--- a/src/qibo/backends/abstract.py
+++ b/src/qibo/backends/abstract.py
@@ -331,6 +331,11 @@ class AbstractBackend(ABC):
         """
         raise_error(NotImplementedError)
 
+    @abstractmethod
+    def squeeze(self, x, axis=None): # pragma: no cover
+        """Removes axis of unit length."""
+        raise_error(NotImplementedError)
+
     @abstractmethod
     def gather(self, x, indices=None, condition=None, axis=0): # pragma: no cover
         """Indexing of one-dimensional tensor."""
diff --git a/src/qibo/backends/numpy.py b/src/qibo/backends/numpy.py
index ad3988baf..991c53197 100644
--- a/src/qibo/backends/numpy.py
+++ b/src/qibo/backends/numpy.py
@@ -167,6 +167,9 @@ class NumpyBackend(abstract.AbstractBackend):
         # Uses numpy backend always (even on Tensorflow)
         return self.np.unique(x, return_counts=return_counts)
 
+    def squeeze(self, x, axis=None):
+        return self.backend.squeeze(x, axis=axis)
+
     def gather(self, x, indices=None, condition=None, axis=0):
         if indices is None:
             if condition is None:
diff --git a/src/qibo/core/gates.py b/src/qibo/core/gates.py
index a38c661e6..1085368b1 100644
--- a/src/qibo/core/gates.py
+++ b/src/qibo/core/gates.py
@@ -93,23 +93,30 @@ class BackendGate(BaseBackendGate):
         if self.is_controlled_by:
             ncontrol = len(self.control_qubits)
             nactive = self.nqubits - ncontrol
+            n = 2 ** ncontrol
             state = K.transpose(state, self.control_cache.order(True))
-            state = K.reshape(state, 2 * (2 ** ncontrol,) + 2 * nactive * (2,))
-
-            updates01 = self.einsum(self.calculation_cache.right0,
-                                    state[:-1, -1], K.conj(self.matrix))
-            updates10 = self.einsum(self.calculation_cache.left0,
-                                    state[-1, :-1],
-                                    self.matrix)
-
-            updates11 = self.einsum(self.calculation_cache.right,
-                                    state[-1, -1], K.conj(self.matrix))
-            updates11 = self.einsum(self.calculation_cache.left,
-                                    updates11, self.matrix)
-
-            updates01 = K.concatenate([state[:-1, :-1], updates01[:, K.newaxis]], axis=1)
-            updates10 = K.concatenate([updates10, updates11[K.newaxis]], axis=0)
-            state = K.concatenate([updates01, updates10[K.newaxis]], axis=0)
+            state = K.reshape(state, 2 * (n,) + 2 * nactive * (2,))
+            state01 = K.gather(state, indices=range(n - 1), axis=0)
+            state01 = K.squeeze(K.gather(state01, indices=[n - 1], axis=1), axis=1)
+            state01 = self.einsum(self.calculation_cache.right0,
+                                  state01, K.conj(self.matrix))
+            state10 = K.gather(state, indices=range(n - 1), axis=1)
+            state10 = K.squeeze(K.gather(state10, indices=[n - 1], axis=0), axis=0)
+            state10 = self.einsum(self.calculation_cache.left0,
+                                  state10, self.matrix)
+
+            state11 = K.squeeze(K.gather(state, indices=[n - 1], axis=0), axis=0)
+            state11 = K.squeeze(K.gather(state11, indices=[n - 1], axis=0), axis=0)
+            state11 = self.einsum(self.calculation_cache.right, state11,
+                                  K.conj(self.matrix))
+            state11 = self.einsum(self.calculation_cache.left,
+                                  state11, self.matrix)
+
+            state00 = K.gather(state, indices=range(n - 1), axis=0)
+            state00 = K.gather(state00, indices=range(n - 1), axis=1)
+            state01 = K.concatenate([state00, state01[:, K.newaxis]], axis=1)
+            state10 = K.concatenate([state10, state11[K.newaxis]], axis=0)
+            state = K.concatenate([state01, state10[K.newaxis]], axis=0)
             state = K.reshape(state, 2 * self.nqubits * (2,))
             state = K.transpose(state, self.control_cache.reverse(True))
         else:


Test patch:
diff --git a/src/qibo/tests_new/test_backends_agreement.py b/src/qibo/tests_new/test_backends_agreement.py
index 339d2cb3c..d67988d91 100644
--- a/src/qibo/tests_new/test_backends_agreement.py
+++ b/src/qibo/tests_new/test_backends_agreement.py
@@ -46,6 +46,7 @@ METHODS = [
     ("eigvalsh", [rand((4, 4))]),
     ("less", [rand(10), rand(10)]),
     ("array_equal", [rand(10), rand(10)]),
+    ("squeeze", {"x": rand((5, 1, 2)), "axis": 1}),
     ("gather_nd", [rand((5, 3)), [0, 1]]),
     ("initial_state", [5, True]),
     ("initial_state", [3, False])