diff options
Diffstat (limited to 'ot/gpu/utils.py')
-rw-r--r-- | ot/gpu/utils.py | 15 |
1 files changed, 9 insertions, 6 deletions
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 |