Repository: qiboteam/qibo

Issue:
Set threads warning messages
I was going to comment this on the qibojit PR but since this is merged I am writing here. I believe the following point is still open, right?

> @stavros11 let me put this info here, I believe we should include a warning message about `set_threads` when if the backend is tensorflow or numpy, in fact, I believe this function does not work as expected for these backends.

I was adding the warnings now, an issue though is that the [`set_threads` method ](https://github.com/qiboteam/qibo/blob/47ad4adc0674ba15578795e7f91c10b4e3823184/src/qibo/backends/__init__.py#L211) that is exposed to the user changes the threads for all created backends. This means that if numpy/tensorflow backends are created but a different backend (qibotf or qibojit) is active, then the numpy/tensorflow warnings will be raised if the user uses `set_threads`. @scarrazza, do you think we should modify `set_threads` to only change the threads of the active backend?

Discussion:

Patch:
diff --git a/src/qibo/backends/__init__.py b/src/qibo/backends/__init__.py
index 825270519..dcd05483d 100644
--- a/src/qibo/backends/__init__.py
+++ b/src/qibo/backends/__init__.py
@@ -105,8 +105,6 @@ class Backend:
             new_backend = self.available_backends.get(name)()
             if self.active_backend is not None:
                 new_backend.set_precision(self.active_backend.precision)
-                if self.active_backend.default_device:
-                    new_backend.set_device(self.active_backend.default_device)
             self.constructed_backends[name] = new_backend
         return self.constructed_backends.get(name)
 
@@ -199,8 +197,9 @@ def set_device(name):
     """
     if not config.ALLOW_SWITCHERS and name != K.default_device:
         log.warning("Device should not be changed after allocating gates.")
+    K.set_device(name)
     for bk in K.constructed_backends.values():
-        if bk.name != "numpy":
+        if bk.name != "numpy" and bk != K.active_backend:
             bk.set_device(name)
 
 
@@ -214,8 +213,7 @@ def set_threads(nthreads):
     Args:
         nthreads (int): number of threads.
     """
-    for bk in K.constructed_backends.values():
-        bk.set_threads(nthreads)
+    K.set_threads(nthreads)
 
 
 def get_threads():
diff --git a/src/qibo/backends/numpy.py b/src/qibo/backends/numpy.py
index 8394e6908..24f707ff5 100644
--- a/src/qibo/backends/numpy.py
+++ b/src/qibo/backends/numpy.py
@@ -25,11 +25,19 @@ class NumpyBackend(abstract.AbstractBackend):
         self.newaxis = np.newaxis
         self.oom_error = MemoryError
         self.optimization = None
+        self.cpu_devices = ["/CPU:0"]
+        self.gpu_devices = []
+        self.default_device = self.cpu_devices[0]
 
     def set_device(self, name):
         log.warning("Numpy does not support device placement. "
                     "Aborting device change.")
 
+    def set_threads(self, nthreads):
+        log.warning("Numpy backend supports only single-thread execution. "
+                    "Cannot change the number of threads.")
+        abstract.AbstractBackend.set_threads(self, nthreads)
+
     def to_numpy(self, x):
         return x
 
@@ -404,7 +412,6 @@ class JITCustomBackend(NumpyBackend): # pragma: no cover
         if "NUMBA_NUM_THREADS" in os.environ: # pragma: no cover
             self.set_threads(int(os.environ.get("NUMBA_NUM_THREADS")))
 
-        # TODO: reconsider device management
         self.cpu_devices = ["/CPU:0"]
         self.gpu_devices = [f"/GPU:{i}" for i in range(ngpu)]
         if self.gpu_devices: # pragma: no cover
@@ -446,7 +453,7 @@ class JITCustomBackend(NumpyBackend): # pragma: no cover
             self.set_engine("numba")
 
     def set_threads(self, nthreads):
-        super().set_threads(nthreads)
+        abstract.AbstractBackend.set_threads(self, nthreads)
         import numba # pylint: disable=E0401
         numba.set_num_threads(nthreads)
 
diff --git a/src/qibo/backends/tensorflow.py b/src/qibo/backends/tensorflow.py
index 5f93d53db..f74062800 100644
--- a/src/qibo/backends/tensorflow.py
+++ b/src/qibo/backends/tensorflow.py
@@ -49,6 +49,14 @@ class TensorflowBackend(numpy.NumpyBackend):
     def set_device(self, name):
         abstract.AbstractBackend.set_device(self, name)
 
+    def set_threads(self, nthreads):
+        log.warning("`set_threads` is not supported by the tensorflow "
+                    "backend. Please use tensorflow's thread setters: "
+                    "`tf.config.threading.set_inter_op_parallelism_threads` "
+                    "or `tf.config.threading.set_intra_op_parallelism_threads` "
+                    "to switch the number of threads.")
+        abstract.AbstractBackend.set_threads(self, nthreads)
+
     def to_numpy(self, x):
         if isinstance(x, self.np.ndarray):
             return x
@@ -216,6 +224,9 @@ class TensorflowCustomBackend(TensorflowBackend):
         if "OMP_NUM_THREADS" in os.environ: # pragma: no cover
             self.set_threads(int(os.environ.get("OMP_NUM_THREADS")))
 
+    def set_threads(self, nthreads):
+        abstract.AbstractBackend.set_threads(self, nthreads)
+
     def initial_state(self, nqubits, is_matrix=False):
         return self.op.initial_state(nqubits, self.dtypes('DTYPECPX'),
                                     is_matrix=is_matrix,


Test patch:
diff --git a/src/qibo/tests/test_backends_init.py b/src/qibo/tests/test_backends_init.py
index 1aa14a9ac..c2f912b51 100644
--- a/src/qibo/tests/test_backends_init.py
+++ b/src/qibo/tests/test_backends_init.py
@@ -103,6 +103,16 @@ def test_set_device(backend, caplog):
         bk.set_device(device)
 
 
+def test_set_threads(backend, caplog):
+    original_threads = backends.get_threads()
+    bkname = backends.get_backend()
+    backends.set_threads(1)
+    if bkname == "numpy" or bkname == "tensorflow":
+        assert "WARNING" in caplog.text
+    assert backends.get_threads() == 1
+    backends.set_threads(original_threads)
+
+
 def test_set_shot_batch_size():
     import qibo
     assert qibo.get_batch_size() == 2 ** 18