add check_number_threads to ot/lp/__init__.py __all__

Author: eloitanguy

ot/lp/__init__.py

@@ -16,6 +16,7 @@
     free_support_barycenter,
     generalized_free_support_barycenter,
 )
+from ..utils import check_number_threads
 
 # import compiled emd
 from .emd_wrap import emd_1d_sorted
@@ -44,4 +45,5 @@
     "semidiscrete_wasserstein2_unif_circle",
     "dmmot_monge_1dgrid_loss",
     "dmmot_monge_1dgrid_optimize",
+    "check_number_threads",
 ]

ot/lp/network_simplex.py

@@ -11,35 +11,11 @@
 import numpy as np
 import warnings
 
-from ..utils import list_to_array
+from ..utils import list_to_array, check_number_threads
 from ..backend import get_backend
 from .emd_wrap import emd_c, check_result
 
 
-def check_number_threads(numThreads):
-    """Checks whether or not the requested number of threads has a valid value.
-
-    Parameters
-    ----------
-    numThreads : int or str
-        The requested number of threads, should either be a strictly positive integer or "max" or None
-
-    Returns
-    -------
-    numThreads : int
-        Corrected number of threads
-    """
-    if (numThreads is None) or (
-        isinstance(numThreads, str) and numThreads.lower() == "max"
-    ):
-        return -1
-    if (not isinstance(numThreads, int)) or numThreads < 1:
-        raise ValueError(
-            'numThreads should either be "max" or a strictly positive integer'
-        )
-    return numThreads
-
-
 def center_ot_dual(alpha0, beta0, a=None, b=None):
     r"""Center dual OT potentials w.r.t. their weights
 

ot/utils.py

@@ -1341,3 +1341,27 @@ def proj_SDP(S, nx=None, vmin=0.0):
         Q = nx.einsum("ijk,ik->ijk", P, w)  # Q[i] = P[i] @ diag(w[i])
         # R[i] = Q[i] @ P[i].T
         return nx.einsum("ijk,ikl->ijl", Q, nx.transpose(P, (0, 2, 1)))
+
+
+def check_number_threads(numThreads):
+    """Checks whether or not the requested number of threads has a valid value.
+
+    Parameters
+    ----------
+    numThreads : int or str
+        The requested number of threads, should either be a strictly positive integer or "max" or None
+
+    Returns
+    -------
+    numThreads : int
+        Corrected number of threads
+    """
+    if (numThreads is None) or (
+        isinstance(numThreads, str) and numThreads.lower() == "max"
+    ):
+        return -1
+    if (not isinstance(numThreads, int)) or numThreads < 1:
+        raise ValueError(
+            'numThreads should either be "max" or a strictly positive integer'
+        )
+    return numThreads