diff options
Diffstat (limited to 'ot/gpu/__init__.py')
| -rw-r--r-- | ot/gpu/__init__.py | 17 |
1 files changed, 9 insertions, 8 deletions
diff --git a/ot/gpu/__init__.py b/ot/gpu/__init__.py index 0187a4f..deda6b1 100644 --- a/ot/gpu/__init__.py +++ b/ot/gpu/__init__.py @@ -9,15 +9,17 @@ 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 +In order to get the best erformances, we recommend to give only cupy arrays to the functions and desactivate the conversion to numpy of the result of the function with parameter ``to_numpy=False``. - - - """ +# Author: Remi Flamary <remi.flamary@unice.fr> +# Leo Gautheron <https://github.com/aje> +# +# License: MIT License + from . import bregman from . import da from .bregman import sinkhorn @@ -27,10 +29,9 @@ from . import utils from .utils import dist, to_gpu, to_np -# Author: Remi Flamary <remi.flamary@unice.fr> -# Leo Gautheron <https://github.com/aje> -# -# License: MIT License + + __all__ = ["utils", "dist", "sinkhorn", "sinkhorn_lpl1_mm", 'bregman', 'da', 'to_gpu', 'to_np'] + |