Repository: qiboteam/qibo

Issue:
Promote partial trace (states and density matrices) from Backend to Quantum Information Module
There are some features in the Backend that could be promoted into some other modules as they are useful in a variety of situations. Partial traces are one of them. But maybe more are also good to promote.

Discussion:

Patch:
diff --git a/doc/source/api-reference/qibo.rst b/doc/source/api-reference/qibo.rst
index afd8c9422..09c88a3fc 100644
--- a/doc/source/api-reference/qibo.rst
+++ b/doc/source/api-reference/qibo.rst
@@ -1930,6 +1930,18 @@ Anticommutator
 .. autofunction:: qibo.quantum_info.anticommutator
 
 
+Partial trace
+"""""""""""""
+
+.. autofunction:: qibo.quantum_info.partial_trace
+
+
+Matrix exponentiation
+"""""""""""""""""""""
+
+.. autofunction:: qibo.quantum_info.matrix_exponentiation
+
+
 Quantum Networks
 ^^^^^^^^^^^^^^^^
 
diff --git a/src/qibo/backends/abstract.py b/src/qibo/backends/abstract.py
index 8e4175b8e..5a14c2402 100644
--- a/src/qibo/backends/abstract.py
+++ b/src/qibo/backends/abstract.py
@@ -66,7 +66,7 @@ class Backend(abc.ABC):
         raise_error(NotImplementedError)
 
     @abc.abstractmethod
-    def issparse(self, x):  # pragma: no cover
+    def is_sparse(self, x):  # pragma: no cover
         """Determine if a given array is a sparse tensor."""
         raise_error(NotImplementedError)
 
@@ -281,16 +281,6 @@ class Backend(abc.ABC):
         """Sample measurement frequencies according to a probability distribution."""
         raise_error(NotImplementedError)
 
-    @abc.abstractmethod
-    def partial_trace(self, state, qubits, nqubits):  # pragma: no cover
-        """Trace out specific qubits of a state vector."""
-        raise_error(NotImplementedError)
-
-    @abc.abstractmethod
-    def partial_trace_density_matrix(self, state, qubits, nqubits):  # pragma: no cover
-        """Trace out specific qubits of a density matrix."""
-        raise_error(NotImplementedError)
-
     @abc.abstractmethod
     def calculate_norm(self, state, order=2):  # pragma: no cover
         """Calculate norm of a state vector.
@@ -337,17 +327,6 @@ class Backend(abc.ABC):
         """Calculate eigenvectors of a matrix."""
         raise_error(NotImplementedError)
 
-    @abc.abstractmethod
-    def calculate_matrix_exp(
-        self, matrix, a, eigenvectors=None, eigenvalues=None
-    ):  # pragma: no cover
-        """Calculate matrix exponential of a matrix.
-
-        If the eigenvectors and eigenvalues are given the matrix diagonalization is
-        used for exponentiation.
-        """
-        raise_error(NotImplementedError)
-
     @abc.abstractmethod
     def calculate_expectation_state(
         self, hamiltonian, state, normalize
@@ -362,6 +341,16 @@ class Backend(abc.ABC):
         """Calculate expectation value of a density matrix given the observable matrix."""
         raise_error(NotImplementedError)
 
+    @abc.abstractmethod
+    def calculate_matrix_exp(
+        self, a, matrix, eigenvectors=None, eigenvalues=None
+    ):  # pragma: no cover
+        """Calculate matrix exponential of a matrix.
+        If the eigenvectors and eigenvalues are given the matrix diagonalization is
+        used for exponentiation.
+        """
+        raise_error(NotImplementedError)
+
     @abc.abstractmethod
     def calculate_hamiltonian_matrix_product(
         self, matrix1, matrix2
@@ -384,7 +373,7 @@ class Backend(abc.ABC):
         self.assert_allclose(value, target, rtol=rtol, atol=atol)
 
     @abc.abstractmethod
-    def test_regressions(self, name):  # pragma: no cover
+    def _test_regressions(self, name):  # pragma: no cover
         """Correct outcomes for tests that involve random numbers.
 
         The outcomes of such tests depend on the backend.
diff --git a/src/qibo/backends/numpy.py b/src/qibo/backends/numpy.py
index ec4329ff0..d69765058 100644
--- a/src/qibo/backends/numpy.py
+++ b/src/qibo/backends/numpy.py
@@ -61,17 +61,17 @@ class NumpyBackend(Backend):
             dtype = self.dtype
         if isinstance(x, self.tensor_types):
             return x.astype(dtype, copy=copy)
-        elif self.issparse(x):
+        elif self.is_sparse(x):
             return x.astype(dtype, copy=copy)
         return np.array(x, dtype=dtype, copy=copy)
 
-    def issparse(self, x):
+    def is_sparse(self, x):
         from scipy import sparse
 
         return sparse.issparse(x)
 
     def to_numpy(self, x):
-        if self.issparse(x):
+        if self.is_sparse(x):
             return x.toarray()
         return x
 
@@ -308,13 +308,16 @@ class NumpyBackend(Backend):
         return self.np.reshape(state, shape)
 
     def reset_error_density_matrix(self, gate, state, nqubits):
-        from qibo.gates import X
+        from qibo.gates import X  # pylint: disable=C0415
+        from qibo.quantum_info.linalg_operations import (  # pylint: disable=C0415
+            partial_trace,
+        )
 
         state = self.cast(state)
         shape = state.shape
         q = gate.target_qubits[0]
         p_0, p_1 = gate.init_kwargs["p_0"], gate.init_kwargs["p_1"]
-        trace = self.partial_trace_density_matrix(state, (q,), nqubits)
+        trace = partial_trace(state, (q,), backend=self)
         trace = self.np.reshape(trace, 2 * (nqubits - 1) * (2,))
         zero = self.zero_density_matrix(1)
         zero = self.np.tensordot(trace, zero, 0)
@@ -332,11 +335,15 @@ class NumpyBackend(Backend):
         return self.np.reshape(state, shape)
 
     def depolarizing_error_density_matrix(self, gate, state, nqubits):
+        from qibo.quantum_info.linalg_operations import (  # pylint: disable=C0415
+            partial_trace,
+        )
+
         state = self.cast(state)
         shape = state.shape
         q = gate.target_qubits
         lam = gate.init_kwargs["lam"]
-        trace = self.partial_trace_density_matrix(state, q, nqubits)
+        trace = partial_trace(state, q, backend=self)
         trace = self.np.reshape(trace, 2 * (nqubits - len(q)) * (2,))
         identity = self.identity_density_matrix(len(q))
         identity = self.np.reshape(identity, 2 * len(q) * (2,))
@@ -687,27 +694,6 @@ class NumpyBackend(Backend):
         noisy_samples = noisy_samples - noiseless_samples * flip_1
         return noisy_samples
 
-    def partial_trace(self, state, qubits, nqubits):
-        state = self.cast(state)
-        state = self.np.reshape(state, nqubits * (2,))
-        axes = 2 * [list(qubits)]
-        rho = self.np.tensordot(state, self.np.conj(state), axes)
-        shape = 2 * (2 ** (nqubits - len(qubits)),)
-        return self.np.reshape(rho, shape)
-
-    def partial_trace_density_matrix(self, state, qubits, nqubits):
-        state = self.cast(state)
-        state = self.np.reshape(state, 2 * nqubits * (2,))
-
-        order = tuple(sorted(qubits))
-        order += tuple(i for i in range(nqubits) if i not in qubits)
-        order += tuple(i + nqubits for i in order)
-        shape = 2 * (2 ** len(qubits), 2 ** (nqubits - len(qubits)))
-
-        state = self.np.transpose(state, order)
-        state = self.np.reshape(state, shape)
-        return self.np.einsum("abac->bc", state)
-
     def calculate_norm(self, state, order=2):
         state = self.cast(state)
         return self.np.linalg.norm(state, order)
@@ -727,7 +713,7 @@ class NumpyBackend(Backend):
         )
 
     def calculate_eigenvalues(self, matrix, k=6, hermitian=True):
-        if self.issparse(matrix):
+        if self.is_sparse(matrix):
             log.warning(
                 "Calculating sparse matrix eigenvectors because "
                 "sparse modules do not provide ``eigvals`` method."
@@ -738,7 +724,7 @@ class NumpyBackend(Backend):
         return np.linalg.eigvals(matrix)
 
     def calculate_eigenvectors(self, matrix, k=6, hermitian=True):
-        if self.issparse(matrix):
+        if self.is_sparse(matrix):
             if k < matrix.shape[0]:
                 from scipy.sparse.linalg import eigsh
 
@@ -749,17 +735,6 @@ class NumpyBackend(Backend):
             return np.linalg.eigh(matrix)
         return np.linalg.eig(matrix)
 
-    def calculate_matrix_exp(self, a, matrix, eigenvectors=None, eigenvalues=None):
-        if eigenvectors is None or self.issparse(matrix):
-            if self.issparse(matrix):
-                from scipy.sparse.linalg import expm
-            else:
-                from scipy.linalg import expm
-            return expm(-1j * a * matrix)
-        expd = self.np.diag(self.np.exp(-1j * a * eigenvalues))
-        ud = self.np.transpose(np.conj(eigenvectors))
-        return self.np.matmul(eigenvectors, self.np.matmul(expd, ud))
-
     def calculate_expectation_state(self, hamiltonian, state, normalize):
         statec = self.np.conj(state)
         hstate = hamiltonian @ state
@@ -775,6 +750,17 @@ class NumpyBackend(Backend):
             ev /= norm
         return ev
 
+    def calculate_matrix_exp(self, a, matrix, eigenvectors=None, eigenvalues=None):
+        if eigenvectors is None or self.is_sparse(matrix):
+            if self.is_sparse(matrix):
+                from scipy.sparse.linalg import expm
+            else:
+                from scipy.linalg import expm
+            return expm(-1j * a * matrix)
+        expd = self.np.diag(self.np.exp(-1j * a * eigenvalues))
+        ud = self.np.transpose(np.conj(eigenvectors))
+        return self.np.matmul(eigenvectors, self.np.matmul(expd, ud))
+
     # TODO: remove this method
     def calculate_hamiltonian_matrix_product(self, matrix1, matrix2):
         return matrix1 @ matrix2
@@ -797,7 +783,7 @@ class NumpyBackend(Backend):
         target = self.to_numpy(target)
         np.testing.assert_allclose(value, target, rtol=rtol, atol=atol)
 
-    def test_regressions(self, name):
+    def _test_regressions(self, name):
         if name == "test_measurementresult_apply_bitflips":
             return [
                 [0, 0, 0, 0, 2, 3, 0, 0, 0, 0],
diff --git a/src/qibo/backends/pytorch.py b/src/qibo/backends/pytorch.py
index 69889f6f4..392aeed40 100644
--- a/src/qibo/backends/pytorch.py
+++ b/src/qibo/backends/pytorch.py
@@ -130,11 +130,11 @@ class PyTorchBackend(NumpyBackend):
 
         return x
 
-    def issparse(self, x):
+    def is_sparse(self, x):
         if isinstance(x, self.np.Tensor):
             return x.is_sparse
 
-        return super().issparse(x)
+        return super().is_sparse(x)
 
     def to_numpy(self, x):
         if isinstance(x, list):
@@ -187,7 +187,7 @@ class PyTorchBackend(NumpyBackend):
         return self.np.linalg.eig(matrix)  # pylint: disable=not-callable
 
     def calculate_matrix_exp(self, a, matrix, eigenvectors=None, eigenvalues=None):
-        if eigenvectors is None or self.issparse(matrix):
+        if eigenvectors is None or self.is_sparse(matrix):
             return self.np.linalg.matrix_exp(  # pylint: disable=not-callable
                 -1j * a * matrix
             )
@@ -195,7 +195,7 @@ class PyTorchBackend(NumpyBackend):
         ud = self.np.conj(eigenvectors).T
         return self.np.matmul(eigenvectors, self.np.matmul(expd, ud))
 
-    def test_regressions(self, name):
+    def _test_regressions(self, name):
         if name == "test_measurementresult_apply_bitflips":
             return [
                 [0, 0, 0, 0, 2, 3, 0, 0, 0, 0],
diff --git a/src/qibo/backends/tensorflow.py b/src/qibo/backends/tensorflow.py
index f45ec05ec..6f4deed49 100644
--- a/src/qibo/backends/tensorflow.py
+++ b/src/qibo/backends/tensorflow.py
@@ -92,7 +92,7 @@ class TensorflowBackend(NumpyBackend):
             return self.tf.identity(x)
         return x
 
-    def issparse(self, x):
+    def is_sparse(self, x):
         return isinstance(x, self.tf.sparse.SparseTensor)
 
     def to_numpy(self, x):
@@ -188,13 +188,12 @@ class TensorflowBackend(NumpyBackend):
         return self.tf.linalg.eig(matrix)
 
     def calculate_matrix_exp(self, a, matrix, eigenvectors=None, eigenvalues=None):
-        if eigenvectors is None or self.issparse(matrix):
+        if eigenvectors is None or self.is_sparse(matrix):
             return self.tf.linalg.expm(-1j * a * matrix)
-        else:
-            return super().calculate_matrix_exp(a, matrix, eigenvectors, eigenvalues)
+        return super().calculate_matrix_exp(a, matrix, eigenvectors, eigenvalues)
 
     def calculate_hamiltonian_matrix_product(self, matrix1, matrix2):
-        if self.issparse(matrix1) or self.issparse(matrix2):
+        if self.is_sparse(matrix1) or self.is_sparse(matrix2):
             raise_error(
                 NotImplementedError,
                 "Multiplication of sparse matrices is not supported with Tensorflow.",
@@ -213,7 +212,7 @@ class TensorflowBackend(NumpyBackend):
                 f"Cannot multiply Hamiltonian with rank-{rank} tensor.",
             )
 
-    def test_regressions(self, name):
+    def _test_regressions(self, name):
         if name == "test_measurementresult_apply_bitflips":
             return [
                 [4, 0, 0, 1, 0, 0, 1, 0, 0, 0],
diff --git a/src/qibo/hamiltonians/hamiltonians.py b/src/qibo/hamiltonians/hamiltonians.py
index 4bdec69d3..bceb1b3b0 100644
--- a/src/qibo/hamiltonians/hamiltonians.py
+++ b/src/qibo/hamiltonians/hamiltonians.py
@@ -31,7 +31,7 @@ class Hamiltonian(AbstractHamiltonian):
 
         if not (
             isinstance(matrix, self.backend.tensor_types)
-            or self.backend.issparse(matrix)
+            or self.backend.is_sparse(matrix)
         ):
             raise_error(
                 TypeError,
@@ -107,10 +107,14 @@ class Hamiltonian(AbstractHamiltonian):
         return self._eigenvectors
 
     def exp(self, a):
+        from qibo.quantum_info.linalg_operations import (  # pylint: disable=C0415
+            matrix_exponentiation,
+        )
+
         if self._exp.get("a") != a:
             self._exp["a"] = a
-            self._exp["result"] = self.backend.calculate_matrix_exp(
-                a, self.matrix, self._eigenvectors, self._eigenvalues
+            self._exp["result"] = matrix_exponentiation(
+                a, self.matrix, self._eigenvectors, self._eigenvalues, self.backend
             )
         return self._exp.get("result")
 
@@ -254,7 +258,7 @@ class Hamiltonian(AbstractHamiltonian):
         if self._eigenvalues is not None:
             if self.backend.np.real(o) >= 0:  # TODO: check for side effects K.qnp
                 r._eigenvalues = o * self._eigenvalues
-            elif not self.backend.issparse(self.matrix):
+            elif not self.backend.is_sparse(self.matrix):
                 axis = (0,) if isinstance(self.backend, PyTorchBackend) else 0
                 r._eigenvalues = o * self.backend.np.flip(self._eigenvalues, axis)
         if self._eigenvectors is not None:
diff --git a/src/qibo/models/dbi/double_bracket.py b/src/qibo/models/dbi/double_bracket.py
index 4c544ea2a..c5b3b1060 100644
--- a/src/qibo/models/dbi/double_bracket.py
+++ b/src/qibo/models/dbi/double_bracket.py
@@ -12,6 +12,7 @@ from qibo.models.dbi.utils_scheduling import (
     polynomial_step,
     simulated_annealing_step,
 )
+from qibo.quantum_info.linalg_operations import commutator, matrix_exponentiation
 
 
 class DoubleBracketGeneratorType(Enum):
@@ -141,43 +142,53 @@ class DoubleBracketIteration:
             mode = self.mode
 
         if mode is DoubleBracketGeneratorType.canonical:
-            operator = self.backend.calculate_matrix_exp(
+            operator = matrix_exponentiation(
                 -1.0j * step,
                 self.commutator(self.diagonal_h_matrix, self.h.matrix),
+                backend=self.backend,
             )
         elif mode is DoubleBracketGeneratorType.single_commutator:
             if d is None:
                 d = self.diagonal_h_matrix
-            operator = self.backend.calculate_matrix_exp(
+            operator = matrix_exponentiation(
                 -1.0j * step,
                 self.commutator(self.backend.cast(d), self.h.matrix),
+                backend=self.backend,
             )
         elif mode is DoubleBracketGeneratorType.group_commutator:
             if d is None:
                 d = self.diagonal_h_matrix
             operator = (
                 self.h.exp(step)
-                @ self.backend.calculate_matrix_exp(-step, d)
+                @ matrix_exponentiation(-step, d, backend=self.backend)
                 @ self.h.exp(-step)
-                @ self.backend.calculate_matrix_exp(step, d)
+                @ matrix_exponentiation(step, d, backend=self.backend)
             )
         elif mode is DoubleBracketGeneratorType.group_commutator_third_order:
             if d is None:
                 d = self.diagonal_h_matrix
             operator = (
                 self.h.exp(-step * (np.sqrt(5) - 1) / 2)
-                @ self.backend.calculate_matrix_exp(-step * (np.sqrt(5) - 1) / 2, d)
+                @ matrix_exponentiation(
+                    -step * (np.sqrt(5) - 1) / 2, d, backend=self.backend
+                )
                 @ self.h.exp(step)
-                @ self.backend.calculate_matrix_exp(step * (np.sqrt(5) + 1) / 2, d)
+                @ matrix_exponentiation(
+                    step * (np.sqrt(5) + 1) / 2, d, backend=self.backend
+                )
                 @ self.h.exp(-step * (3 - np.sqrt(5)) / 2)
-                @ self.backend.calculate_matrix_exp(-step, d)
+                @ matrix_exponentiation(-step, d, backend=self.backend)
             )
             operator = (
-                self.backend.calculate_matrix_exp(step, d)
+                matrix_exponentiation(step, d, backend=self.backend)
                 @ self.h.exp(step * (3 - np.sqrt(5)) / 2)
-                @ self.backend.calculate_matrix_exp(-step * (np.sqrt(5) + 1) / 2, d)
+                @ matrix_exponentiation(
+                    -step * (np.sqrt(5) + 1) / 2, d, backend=self.backend
+                )
                 @ self.h.exp(-step)
-                @ self.backend.calculate_matrix_exp(step * (np.sqrt(5) - 1) / 2, d)
+                @ matrix_exponentiation(
+                    step * (np.sqrt(5) - 1) / 2, d, backend=self.backend
+                )
                 @ self.h.exp(step * (np.sqrt(5) - 1) / 2)
             )
         else:
@@ -190,7 +201,7 @@ class DoubleBracketIteration:
     @staticmethod
     def commutator(a, b):
         """Compute commutator between two arrays."""
-        return a @ b - b @ a
+        return commutator(a, b)
 
     @property
     def diagonal_h_matrix(self):
diff --git a/src/qibo/quantum_info/__init__.py b/src/qibo/quantum_info/__init__.py
index 6756df73a..9835cb962 100644
--- a/src/qibo/quantum_info/__init__.py
+++ b/src/qibo/quantum_info/__init__.py
@@ -2,8 +2,8 @@ from qibo.quantum_info.basis import *
 from qibo.quantum_info.clifford import *
 from qibo.quantum_info.entanglement import *
 from qibo.quantum_info.entropies import *
+from qibo.quantum_info.linalg_operations import *
 from qibo.quantum_info.metrics import *
-from qibo.quantum_info.operations import *
 from qibo.quantum_info.quantum_networks import *
 from qibo.quantum_info.random_ensembles import *
 from qibo.quantum_info.superoperator_transformations import *
diff --git a/src/qibo/quantum_info/entanglement.py b/src/qibo/quantum_info/entanglement.py
index e1f124bd6..23c2f7d3f 100644
--- a/src/qibo/quantum_info/entanglement.py
+++ b/src/qibo/quantum_info/entanglement.py
@@ -4,6 +4,7 @@ import numpy as np
 
 from qibo.backends import _check_backend
 from qibo.config import PRECISION_TOL, raise_error
+from qibo.quantum_info.linalg_operations import partial_trace
 from qibo.quantum_info.metrics import fidelity, purity
 
 
@@ -47,8 +48,6 @@ def concurrence(state, bipartition, check_purity: bool = True, backend=None):
             f"check_purity must be type bool, but it is type {type(check_purity)}.",
         )
 
-    nqubits = int(np.log2(state.shape[0]))
-
     if check_purity is True:
         purity_total_system = purity(state, backend=backend)
 
@@ -59,11 +58,7 @@ def concurrence(state, bipartition, check_purity: bool = True, backend=None):
                 "concurrence only implemented for pure quantum states.",
             )
 
-    reduced_density_matrix = (
-        backend.partial_trace(state, bipartition, nqubits)
-        if len(state.shape) == 1
-        else backend.partial_trace_density_matrix(state, bipartition, nqubits)
-    )
+    reduced_density_matrix = partial_trace(state, bipartition, backend=backend)
 
     purity_reduced = purity(reduced_density_matrix, backend=backend)
     if purity_reduced - 1.0 > 0.0:
@@ -227,7 +222,7 @@ def meyer_wallach_entanglement(circuit, backend=None):
         trace_q = list(range(nqubits))
         trace_q.pop(j)
 
-        rho_r = backend.partial_trace_density_matrix(rho, trace_q, nqubits)
+        rho_r = partial_trace(rho, trace_q, backend=backend)
 
         trace = purity(rho_r, backend=backend)
 
diff --git a/src/qibo/quantum_info/entropies.py b/src/qibo/quantum_info/entropies.py
index 20ff3e2af..05bf3f823 100644
--- a/src/qibo/quantum_info/entropies.py
+++ b/src/qibo/quantum_info/entropies.py
@@ -8,6 +8,7 @@ from scipy.linalg import fractional_matrix_power
 from qibo.backends import _check_backend
 from qibo.backends.pytorch import PyTorchBackend
 from qibo.config import PRECISION_TOL, raise_error
+from qibo.quantum_info.linalg_operations import partial_trace
 from qibo.quantum_info.metrics import _check_hermitian_or_not_gpu, purity
 
 
@@ -869,13 +870,7 @@ def entanglement_entropy(
             f"state must have dims either (k,) or (k,k), but have dims {state.shape}.",
         )
 
-    nqubits = int(np.log2(state.shape[0]))
-
-    reduced_density_matrix = (
-        backend.partial_trace(state, bipartition, nqubits)
-        if len(state.shape) == 1
-        else backend.partial_trace_density_matrix(state, bipartition, nqubits)
-    )
+    reduced_density_matrix = partial_trace(state, bipartition, backend=backend)
 
     entropy_entanglement = von_neumann_entropy(
         reduced_density_matrix,
diff --git a/src/qibo/quantum_info/linalg_operations.py b/src/qibo/quantum_info/linalg_operations.py
new file mode 100644
index 000000000..ae87ba4e1
--- /dev/null
+++ b/src/qibo/quantum_info/linalg_operations.py
@@ -0,0 +1,198 @@
+"""Module with common linear algebra operations for quantum information."""
+
+import math
+from typing import List, Tuple, Union
+
+from qibo.backends import _check_backend
+from qibo.config import raise_error
+
+
+def commutator(operator_1, operator_2):
+    """Returns the commutator of ``operator_1`` and ``operator_2``.
+
+    The commutator of two matrices :math:`A` and :math:`B` is given by
+
+    .. math::
+        [A, B] = A \\, B - B \\, A \\,.
+
+    Args:
+        operator_1 (ndarray): First operator.
+        operator_2 (ndarray): Second operator.
+
+    Returns:
+        ndarray: Commutator of ``operator_1`` and ``operator_2``.
+    """
+    if (
+        (len(operator_1.shape) >= 3)
+        or (len(operator_1) == 0)
+        or (len(operator_1.shape) == 2 and operator_1.shape[0] != operator_1.shape[1])
+    ):
+        raise_error(
+            TypeError,
+            f"``operator_1`` must have shape (k,k), but have shape {operator_1.shape}.",
+        )
+
+    if (
+        (len(operator_2.shape) >= 3)
+        or (len(operator_2) == 0)
+        or (len(operator_2.shape) == 2 and operator_2.shape[0] != operator_2.shape[1])
+    ):
+        raise_error(
+            TypeError,
+            f"``operator_2`` must have shape (k,k), but have shape {operator_2.shape}.",
+        )
+
+    if operator_1.shape != operator_2.shape:
+        raise_error(
+            TypeError,
+            "``operator_1`` and ``operator_2`` must have the same shape, "
+            + f"but {operator_1.shape} != {operator_2.shape}",
+        )
+
+    return operator_1 @ operator_2 - operator_2 @ operator_1
+
+
+def anticommutator(operator_1, operator_2):
+    """Returns the anticommutator of ``operator_1`` and ``operator_2``.
+
+    The anticommutator of two matrices :math:`A` and :math:`B` is given by
+
+    .. math::
+        \\{A, B\\} = A \\, B + B \\, A \\,.
+
+    Args:
+        operator_1 (ndarray): First operator.
+        operator_2 (ndarray): Second operator.
+
+    Returns:
+        ndarray: Anticommutator of ``operator_1`` and ``operator_2``.
+    """
+    if (
+        (len(operator_1.shape) >= 3)
+        or (len(operator_1) == 0)
+        or (len(operator_1.shape) == 2 and operator_1.shape[0] != operator_1.shape[1])
+    ):
+        raise_error(
+            TypeError,
+            f"``operator_1`` must have shape (k,k), but have shape {operator_1.shape}.",
+        )
+
+    if (
+        (len(operator_2.shape) >= 3)
+        or (len(operator_2) == 0)
+        or (len(operator_2.shape) == 2 and operator_2.shape[0] != operator_2.shape[1])
+    ):
+        raise_error(
+            TypeError,
+            f"``operator_2`` must have shape (k,k), but have shape {operator_2.shape}.",
+        )
+
+    if operator_1.shape != operator_2.shape:
+        raise_error(
+            TypeError,
+            "``operator_1`` and ``operator_2`` must have the same shape, "
+            + f"but {operator_1.shape} != {operator_2.shape}",
+        )
+
+    return operator_1 @ operator_2 + operator_2 @ operator_1
+
+
+def partial_trace(state, traced_qubits: Union[List[int], Tuple[int]], backend=None):
+    """Returns the density matrix resulting from tracing out ``traced_qubits`` from ``state``.
+
+    Total number of qubits is inferred by the shape of ``state``.
+
+    Args:
+        state (ndarray): density matrix or statevector.
+        traced_qubits (Union[List[int], Tuple[int]]): indices of qubits to be traced out.
+        backend (:class:`qibo.backends.abstract.Backend`, optional): backend
+            to be used in the execution. If ``None``, it uses
+            :class:`qibo.backends.GlobalBackend`. Defaults to ``None``.
+
+    Returns:
+        ndarray: Density matrix of the remaining qubit(s).
+    """
+    if (
+        (len(state.shape) >= 3)
+        or (len(state) == 0)
+        or (len(state.shape) == 2 and state.shape[0] != state.shape[1])
+    ):
+        raise_error(
+            TypeError,
+            f"``state`` must have dims either (k,) or (k,k), but have dims {state.shape}.",
+        )
+
+    backend = _check_backend(backend)
+
+    state = backend.cast(state, dtype=state.dtype)
+    nqubits = math.log2(state.shape[0])
+
+    if not nqubits.is_integer():
+        raise_error(
+            ValueError,
+            "dimension(s) of ``state`` must be a power of 2, "
+            + f"but it is {2**nqubits}.",
+        )
+
+    nqubits = int(nqubits)
+
+    statevector = bool(len(state.shape) == 1)
+
+    factor = 1 if statevector else 2
+    state = backend.np.reshape(state, factor * nqubits * (2,))
+
+    if statevector:
+        axes = 2 * [list(traced_qubits)]
+        rho = backend.np.tensordot(state, backend.np.conj(state), axes)
+        shape = 2 * (2 ** (nqubits - len(traced_qubits)),)
+
+        return backend.np.reshape(rho, shape)
+
+    order = tuple(sorted(traced_qubits))
+    order += tuple(i for i in range(nqubits) if i not in traced_qubits)
+    order += tuple(i + nqubits for i in order)
+    shape = 2 * (2 ** len(traced_qubits), 2 ** (nqubits - len(traced_qubits)))
+
+    state = backend.np.transpose(state, order)
+    state = backend.np.reshape(state, shape)
+
+    return backend.np.einsum("abac->bc", state)
+
+
+def matrix_exponentiation(
+    phase: Union[float, complex],
+    matrix,
+    eigenvectors=None,
+    eigenvalues=None,
+    backend=None,
+):
+    """Calculates the exponential of a matrix.
+
+    Given a ``matrix`` :math:`H` and a ``phase`` :math:`\\theta`,
+    it returns the exponential of the form
+
+    .. math::
+        \\exp\\left(-i \\, \\theta \\, H \\right) \\, .
+
+    If the ``eigenvectors`` and ``eigenvalues`` are given, the matrix diagonalization
+    is used for the exponentiation.
+
+    Args:
+        phase (float or complex): phase that multiplies the matrix.
+        matrix (ndarray): matrix to be exponentiated.
+        eigenvectors (ndarray, optional): _if not ``None``, eigenvectors are used
+            to calculate ``matrix`` exponentiation as part of diagonalization.
+            Must be used together with ``eigenvalues``. Defaults to ``None``.
+        eigenvalues (ndarray, optional): if not ``None``, eigenvalues are used
+            to calculate ``matrix`` exponentiation as part of diagonalization.
+            Must be used together with ``eigenvectors``. Defaults to ``None``.
+        backend (:class:`qibo.backends.abstract.Backend`, optional): backend
+            to be used in the execution. If ``None``, it uses
+            :class:`qibo.backends.GlobalBackend`. Defaults to ``None``.
+
+    Returns:
+        ndarray: matrix exponential of :math:`-i \\, \\theta \\, H`.
+    """
+    backend = _check_backend(backend)
+
+    return backend.calculate_matrix_exp(phase, matrix, eigenvectors, eigenvalues)
diff --git a/src/qibo/quantum_info/operations.py b/src/qibo/quantum_info/operations.py
deleted file mode 100644
index 29e4ea08f..000000000
--- a/src/qibo/quantum_info/operations.py
+++ /dev/null
@@ -1,93 +0,0 @@
-"""Module with the most common linear algebra operations for quantum information."""
-
-from qibo.config import raise_error
-
-
-def commutator(operator_1, operator_2):
-    """Returns the commutator of ``operator_1`` and ``operator_2``.
-
-    The commutator of two matrices :math:`A` and :math:`B` is given by
-
-    .. math::
-        [A, B] = A \\, B - B \\, A \\,.
-
-    Args:
-        operator_1 (ndarray): First operator.
-        operator_2 (ndarray): Second operator.
-
-    Returns:
-        ndarray: Commutator of ``operator_1`` and ``operator_2``.
-    """
-    if (
-        (len(operator_1.shape) >= 3)
-        or (len(operator_1) == 0)
-        or (len(operator_1.shape) == 2 and operator_1.shape[0] != operator_1.shape[1])
-    ):
-        raise_error(
-            TypeError,
-            f"``operator_1`` must have shape (k,k), but have shape {operator_1.shape}.",
-        )
-
-    if (
-        (len(operator_2.shape) >= 3)
-        or (len(operator_2) == 0)
-        or (len(operator_2.shape) == 2 and operator_2.shape[0] != operator_2.shape[1])
-    ):
-        raise_error(
-            TypeError,
-            f"``operator_2`` must have shape (k,k), but have shape {operator_2.shape}.",
-        )
-
-    if operator_1.shape != operator_2.shape:
-        raise_error(
-            TypeError,
-            "``operator_1`` and ``operator_2`` must have the same shape, "
-            + f"but {operator_1.shape} != {operator_2.shape}",
-        )
-
-    return operator_1 @ operator_2 - operator_2 @ operator_1
-
-
-def anticommutator(operator_1, operator_2):
-    """Returns the anticommutator of ``operator_1`` and ``operator_2``.
-
-    The anticommutator of two matrices :math:`A` and :math:`B` is given by
-
-    .. math::
-        \\{A, B\\} = A \\, B + B \\, A \\,.
-
-    Args:
-        operator_1 (ndarray): First operator.
-        operator_2 (ndarray): Second operator.
-
-    Returns:
-        ndarray: Anticommutator of ``operator_1`` and ``operator_2``.
-    """
-    if (
-        (len(operator_1.shape) >= 3)
-        or (len(operator_1) == 0)
-        or (len(operator_1.shape) == 2 and operator_1.shape[0] != operator_1.shape[1])
-    ):
-        raise_error(
-            TypeError,
-            f"``operator_1`` must have shape (k,k), but have shape {operator_1.shape}.",
-        )
-
-    if (
-        (len(operator_2.shape) >= 3)
-        or (len(operator_2) == 0)
-        or (len(operator_2.shape) == 2 and operator_2.shape[0] != operator_2.shape[1])
-    ):
-        raise_error(
-            TypeError,
-            f"``operator_2`` must have shape (k,k), but have shape {operator_2.shape}.",
-        )
-
-    if operator_1.shape != operator_2.shape:
-        raise_error(
-            TypeError,
-            "``operator_1`` and ``operator_2`` must have the same shape, "
-            + f"but {operator_1.shape} != {operator_2.shape}",
-        )
-
-    return operator_1 @ operator_2 + operator_2 @ operator_1


Test patch:
diff --git a/tests/test_callbacks.py b/tests/test_callbacks.py
index a551245b1..b599a904f 100644
--- a/tests/test_callbacks.py
+++ b/tests/test_callbacks.py
@@ -3,7 +3,7 @@
 import numpy as np
 import pytest
 
-from qibo import callbacks, gates
+from qibo import callbacks, gates, hamiltonians
 
 # Absolute testing tolerance for the cases of zero entanglement entropy
 from qibo.config import PRECISION_TOL
@@ -353,14 +353,9 @@ def test_energy(backend, density_matrix):
 @pytest.mark.parametrize("dense", [False, True])
 @pytest.mark.parametrize("check_degenerate", [False, True])
 def test_gap(backend, dense, check_degenerate):
-    from qibo import hamiltonians
-
     h0 = hamiltonians.X(4, dense=dense, backend=backend)
-    if check_degenerate:
-        # use h=0 to make this Hamiltonian degenerate
-        h1 = hamiltonians.TFIM(4, h=0, dense=dense, backend=backend)
-    else:
-        h1 = hamiltonians.TFIM(4, h=1, dense=dense, backend=backend)
+    h = 0 if check_degenerate else 1
+    h1 = hamiltonians.TFIM(4, h=h, dense=dense, backend=backend)
 
     ham = lambda t: (1 - t) * h0.matrix + t * h1.matrix
     targets = {"ground": [], "excited": [], "gap": []}
diff --git a/tests/test_gates_channels.py b/tests/test_gates_channels.py
index def0cd9a3..a1a797454 100644
--- a/tests/test_gates_channels.py
+++ b/tests/test_gates_channels.py
@@ -6,6 +6,7 @@ import pytest
 from qibo import gates, matrices
 from qibo.config import PRECISION_TOL
 from qibo.quantum_info import (
+    partial_trace,
     random_density_matrix,
     random_statevector,
     random_stochastic_matrix,
@@ -257,14 +258,14 @@ def test_depolarizing_channel_errors():
 
 def test_depolarizing_channel(backend):
     """"""
-    initial_state = random_density_matrix(2**3, backend=backend)
     lam = 0.3
-    initial_state_r = backend.partial_trace_density_matrix(initial_state, (2,), 3)
+    initial_state = random_density_matrix(2**3, backend=backend)
+    initial_state_r = partial_trace(initial_state, (2,), backend=backend)
     channel = gates.DepolarizingChannel((0, 1), lam)
     final_state = channel.apply_density_matrix(
         backend, backend.np.copy(initial_state), 3
     )
-    final_state_r = backend.partial_trace_density_matrix(final_state, (2,), 3)
+    final_state_r = partial_trace(final_state, (2,), backend=backend)
     target_state_r = (1 - lam) * initial_state_r + lam * backend.cast(
         np.identity(4)
     ) / 4
@@ -360,9 +361,7 @@ def test_thermal_relaxation_channel(backend, t_1, t_2, time, excpop):
         m_z = backend.cast(m_z, dtype=m_z.dtype)
         z_rho = m_z @ initial_state @ m_z
 
-        trace = backend.to_numpy(
-            backend.partial_trace_density_matrix(initial_state, (0,), 3)
-        )
+        trace = backend.to_numpy(partial_trace(initial_state, (0,), backend=backend))
         trace = np.reshape(trace, 4 * (2,))
         zeros = np.tensordot(
             trace, np.array([[1, 0], [0, 0]], dtype=trace.dtype), axes=0
@@ -437,9 +436,7 @@ def test_reset_channel(backend):
         gate, backend.np.copy(initial_state), 3
     )
 
-    trace = backend.to_numpy(
-        backend.partial_trace_density_matrix(initial_state, (0,), 3)
-    )
+    trace = backend.to_numpy(partial_trace(initial_state, (0,), backend=backend))
     trace = np.reshape(trace, 4 * (2,))
 
     zeros = np.tensordot(trace, np.array([[1, 0], [0, 0]], dtype=trace.dtype), axes=0)
diff --git a/tests/test_measurements_probabilistic.py b/tests/test_measurements_probabilistic.py
index 13381865f..742c4838b 100644
--- a/tests/test_measurements_probabilistic.py
+++ b/tests/test_measurements_probabilistic.py
@@ -25,7 +25,7 @@ def test_probabilistic_measurement(backend, accelerators, use_samples):
         _ = result.samples()
 
     # update reference values based on backend and device
-    decimal_frequencies = backend.test_regressions("test_probabilistic_measurement")
+    decimal_frequencies = backend._test_regressions("test_probabilistic_measurement")
     assert sum(result.frequencies().values()) == 1000
     assert_result(backend, result, decimal_frequencies=decimal_frequencies)
 
@@ -64,7 +64,7 @@ def test_unbalanced_probabilistic_measurement(backend, use_samples):
         # otherwise it uses the frequency-only calculation
         _ = result.samples()
     # update reference values based on backend and device
-    decimal_frequencies = backend.test_regressions(
+    decimal_frequencies = backend._test_regressions(
         "test_unbalanced_probabilistic_measurement"
     )
 
@@ -115,7 +115,7 @@ def test_post_measurement_bitflips_on_circuit(backend, accelerators, i, probs):
     c.add(gates.M(3, p0=probs[2]))
     result = backend.execute_circuit(c, nshots=30)
     freqs = result.frequencies(binary=False)
-    targets = backend.test_regressions("test_post_measurement_bitflips_on_circuit")
+    targets = backend._test_regressions("test_post_measurement_bitflips_on_circuit")
     assert freqs == targets[i]
 
 
@@ -153,6 +153,6 @@ def test_measurementresult_apply_bitflips(backend, i, p0, p1):
     result._samples = backend.cast(np.zeros((10, 3)), dtype="int32")
     backend.set_seed(123)
     noisy_samples = result.apply_bitflips(p0, p1)
-    targets = backend.test_regressions("test_measurementresult_apply_bitflips")
+    targets = backend._test_regressions("test_measurementresult_apply_bitflips")
     noisy_samples = backend.samples_to_decimal(noisy_samples, 3)
     backend.assert_allclose(noisy_samples, targets[i])
diff --git a/tests/test_quantum_info_operations.py b/tests/test_quantum_info_operations.py
index 0dae19ac7..057b69e6e 100644
--- a/tests/test_quantum_info_operations.py
+++ b/tests/test_quantum_info_operations.py
@@ -1,8 +1,13 @@
 import numpy as np
 import pytest
 
-from qibo import matrices
-from qibo.quantum_info.operations import anticommutator, commutator
+from qibo import Circuit, gates, matrices
+from qibo.quantum_info.linalg_operations import (
+    anticommutator,
+    commutator,
+    partial_trace,
+)
+from qibo.quantum_info.random_ensembles import random_density_matrix, random_statevector
 
 
 def test_commutator(backend):
@@ -75,3 +80,32 @@ def test_anticommutator(backend):
 
     anticomm = anticommutator(X, Z)
     backend.assert_allclose(anticomm, 0.0)
+
+
+@pytest.mark.parametrize("density_matrix", [False, True])
+def test_partial_trace(backend, density_matrix):
+    with pytest.raises(TypeError):
+        state = np.random.rand(2, 2, 2).astype(complex)
+        state += 1j * np.random.rand(2, 2, 2)
+        state = backend.cast(state, dtype=state.dtype)
+        test = partial_trace(state, 1, backend=backend)
+    with pytest.raises(ValueError):
+        state = (
+            random_density_matrix(5, backend=backend)
+            if density_matrix
+            else random_statevector(5, backend=backend)
+        )
+        test = partial_trace(state, 1, backend=backend)
+
+    nqubits = 4
+
+    circuit = Circuit(nqubits, density_matrix=density_matrix)
+    circuit.add(gates.H(0))
+    circuit.add(gates.CNOT(0, qubit + 1) for qubit in range(1, nqubits - 1))
+    state = backend.execute_circuit(circuit).state()
+
+    traced = partial_trace(state, (1, 2, 3), backend=backend)
+
+    Id = backend.identity_density_matrix(1, normalize=True)
+
+    backend.assert_allclose(traced, Id)
diff --git a/tests/test_transpiler_unitary_decompositions.py b/tests/test_transpiler_unitary_decompositions.py
index 69e1ffd6d..df641208b 100644
--- a/tests/test_transpiler_unitary_decompositions.py
+++ b/tests/test_transpiler_unitary_decompositions.py
@@ -4,6 +4,7 @@ from scipy.linalg import expm
 
 from qibo import Circuit, gates, matrices
 from qibo.config import PRECISION_TOL
+from qibo.quantum_info.linalg_operations import partial_trace
 from qibo.quantum_info.metrics import purity
 from qibo.quantum_info.random_ensembles import random_unitary
 from qibo.transpiler.unitary_decompositions import (
@@ -68,7 +69,7 @@ def test_eigenbasis_entanglement(backend, seed):
         eigvals = backend.cast(eigvals, dtype=eigvals.dtype)
         backend.assert_allclose(backend.np.abs(eigvals), np.ones(4))
         for state in states.T:
-            state = backend.partial_trace(state, [1], 2)
+            state = partial_trace(state, [1], backend=backend)
             backend.assert_allclose(purity(state, backend=backend), 0.5)