sort operation

[DavidLapous]

Hi, and many thanks for your work ! I'm trying to optimize some part of my code using pykeops; and one piece that seems pykeops-friendly to me is computing sliced wasserstein kernels, as there is a bunch distance, scalar products and norms to compute there. However, to compute (1d)wasserstein on we need to be able to sort w.r.t. one axis (as the 1d wasserstein distance can be efficiently computed with sort+norm). I basically want to do the following computations to compute a sliced wassersein distance matrix

• For each pair of point measure $\mu$ and $\nu$, i.e., point clouds in $\mathbb R^D$ with the same mass; i.e. of shape num_pts, D
• compute the scalar products of $\mu$ and $\nu$ against a bunch of angles $\theta \in \mathbb S^{n-1}$, this gives us the orthogonal projections of $\mu$ and $\nu$ over the lines $\lbrace{ t\cdot \theta : t \in \mathbb R \rbrace}$
• compute (1d) wasserstein distances with sort + norm (over the num_pts axis)
• take the mean of all slices (over the $\theta$ axis) to get an approximation of the sliced wasserstein distance. and create a kernel by exponentiating it afterwards.

My issue here is that I cannot find a sort operation on lazytensors to compute the wasserstein 1 distance on the "num_points" axis. Is there such a "sort along one axis" operation ? or is this kind of computation out of scope of this library ? I didn't find a way to do it with geomloss either.

[jeanfeydy]

Hi David,

Thanks for your comment! This operation is out of the scope of KeOps, since NumPy/PyTorch already provide an optimal workflow:

• Compute the dot products with a matrix multiply.
• Sort the scores using torch.sort (that should rely on a good CUDA kernel, hopefully). This is a non-trivial operation to implement well, especially on GPU.
• Compute the sliced Wasserstein distance as a simple Lp norm between the sorted arrays.

This is the whole point of sliced metrics: you don't need fancy code to compute them efficiently :-)

Best regards, Jean

[DavidLapous]

Thanks Jean for your answer! I'll stick to my current implementation then ;) -- David