update documentation

4 files changed, 43 insertions, 9 deletions

diff --git a/ot/gpu/__init__.py b/ot/gpu/__init__.py
index de4825d..9de2c40 100644
--- a/ot/gpu/__init__.py
+++ b/ot/gpu/__init__.py
@@ -1,8 +1,28 @@
 # -*- coding: utf-8 -*-
+"""
+
+
+This module implement GPU ilmplementation for several OT solvers and utility 
+functions. The GPU backend in handled by `cupy 
+<https://cupy.chainer.org/>`_.
+
+By default, the functions in this module accept and return numpy arrays 
+in order to proide drop-in replacement for the other POT function but
+the transfer between CPU en GPU comes with a significant overhead.
+
+In order to get the best erformances, we recommend to given only cupy 
+arrays to the functions and desactivate the conversion to numpy of the 
+result of the function with parameter ``to_numpy=False``.
+
+
+
+
+"""
 
 from . import bregman
 from . import da
 from .bregman import sinkhorn
+from .da 
 
 from . import utils
 from .utils import dist, to_gpu, to_np
@@ -13,4 +33,4 @@ from .utils import dist, to_gpu, to_np
 #
 # License: MIT License
 
-__all__ = ["utils", "dist", "sinkhorn"]
+__all__ = ["utils", "dist", "sinkhorn", 'bregman', 'da', 'to_gpu', 'to_np']
diff --git a/ot/gpu/bregman.py b/ot/gpu/bregman.py
index 600ead4..978b307 100644
--- a/ot/gpu/bregman.py
+++ b/ot/gpu/bregman.py
@@ -16,7 +16,10 @@ from . import utils
 def sinkhorn_knopp(a, b, M, reg, numItermax=1000, stopThr=1e-9,
                    verbose=False, log=False, to_numpy=True, **kwargs):
     """
-    Solve the entropic regularization optimal transport problem and return the OT matrix
+    Solve the entropic regularization optimal transport on GPU
+
+    If the input matrix are in numpy format, they will be uploaded to the
+    GPU first which can incur significant time overhead.
 
     The function solves the following optimization problem:
 
@@ -56,6 +59,8 @@ def sinkhorn_knopp(a, b, M, reg, numItermax=1000, stopThr=1e-9,
         Print information along iterations
     log : bool, optional
         record log if True
+    to_numpy : boolean, optional (default True)
+        If true convert back the GPU array result to numpy format.
 
 
     Returns
diff --git a/ot/gpu/da.py b/ot/gpu/da.py
index 8c63870..6aba29c 100644
--- a/ot/gpu/da.py
+++ b/ot/gpu/da.py
@@ -22,7 +22,11 @@ def sinkhorn_lpl1_mm(a, labels_a, b, M, reg, eta=0.1, numItermax=10,
                      log=False, to_numpy=True):
     """
     Solve the entropic regularization optimal transport problem with nonconvex
-    group lasso regularization
+    group lasso regularization on GPU
+
+    If the input matrix are in numpy format, they will be uploaded to the
+    GPU first which can incur significant time overhead.
+
 
     The function solves the following optimization problem:
 
@@ -74,6 +78,8 @@ def sinkhorn_lpl1_mm(a, labels_a, b, M, reg, eta=0.1, numItermax=10,
         Print information along iterations
     log : bool, optional
         record log if True
+    to_numpy : boolean, optional (default True)
+        If true convert back the GPU array result to numpy format.
 
 
     Returns
diff --git a/ot/gpu/utils.py b/ot/gpu/utils.py
index d349a6d..41e168a 100644
--- a/ot/gpu/utils.py
+++ b/ot/gpu/utils.py
@@ -16,19 +16,22 @@ import cupy as cp  # cp used for cupy specific operations
 def euclidean_distances(a, b, squared=False, to_numpy=True):
     """
     Compute the pairwise euclidean distance between matrices a and b.
-     Parameters
+
+    If the input matrix are in numpy format, they will be uploaded to the
+    GPU first which can incur significant time overhead.
+
+    Parameters
     ----------
     a : np.ndarray (n, f)
         first matrix
     b : np.ndarray (m, f)
         second matrix
-    gpu : boolean, optional (default False)
-        if True and the module cupy is available, the computation is done
-        on the GPU and the type of the matrix returned is cupy.ndarray.
-        Otherwise, compute on the CPU and returns np.ndarray.
+    to_numpy : boolean, optional (default True)
+        If true convert back the GPU array result to numpy format.
     squared : boolean, optional (default False)
         if True, return squared euclidean distance matrix
-     Returns
+
+    Returns
     -------
     c : (n x m) np.ndarray or cupy.ndarray
         pairwise euclidean distance distance matrix