feat: moving numba threads setting out of the init

src/qibojit/backends/cpu.py

@@ -26,10 +26,11 @@
 class NumbaBackend(NumpyBackend):
     def __init__(self):
         super().__init__()
-        import sys
+        # import sys
 
-        import psutil
+        # import psutil
         from numba import __version__ as numba_version
+        from numba import get_num_threads
 
         from qibojit import __version__ as qibojit_version
         from qibojit.custom_operators import gates, ops
@@ -64,23 +65,27 @@ def __init__(self):
             4: self.gates.apply_four_qubit_gate_kernel,
             5: self.gates.apply_five_qubit_gate_kernel,
         }
+        """
         if sys.platform == "darwin":  # pragma: no cover
             self.set_threads(psutil.cpu_count(logical=False))
         else:
             self.set_threads(len(psutil.Process().cpu_affinity()))
+        """
+        self.nthreads = get_num_threads()
 
     def set_precision(self, precision):
         if precision != self.precision:
             super().set_precision(precision)
             if self.custom_matrices:
                 self.custom_matrices = CustomMatrices(self.dtype)
 
+    """
     def set_threads(self, nthreads):
         import numba
 
         numba.set_num_threads(nthreads)
         self.nthreads = nthreads
-
+    """
     # def cast(self, x, dtype=None, copy=False): Inherited from ``NumpyBackend``
 
     # def to_numpy(self, x): Inherited from ``NumpyBackend``

src/qibojit/custom_operators/ops.py

@@ -1,5 +1,19 @@
+import os
+import sys
+
 import numpy as np
-from numba import njit, prange
+import psutil
+from numba import njit, prange, set_num_threads
+
+NTHREADS = (
+    psutil.cpu_count(logical=False)
+    if sys.platform == "darwin"
+    else len(psutil.Process().cpu_affinity())
+)
+MAX_THREADS = os.environ.get("NUMBA_NUM_THREADS")
+if MAX_THREADS is not None:
+    NTHREADS = min(NTHREADS, MAX_THREADS)
+set_num_threads(NTHREADS)
 
 
 @njit(