- ott.tools.soft_sort.sort_with(inputs, criterion, topk=- 1, **kwargs)#
Sort a multidimensional array according to a real valued criterion.
batchvectors of dimension dim, to which, for each, a real value
criterionis associated, this function produces
topkis set to -1, as by default) composite vectors of size
dimthat will be convex combinations of all vectors, ranked from smallest to largest criterion. Composite vectors with the largest indices will contain convex combinations of those vectors with highest criterion, vectors with smaller indices will contain combinations of vectors with smaller criterion.
ndarray) – the inputs as a jnp.ndarray[batch, dim].
ndarray) – the values according to which to sort the inputs. It has shape [batch, 1].
int) – The number of outputs to keep.
Any) – keyword arguments passed on to lower level functions. Of interest to the user are
squashing_fun, which will redistribute the values in
inputsto lie in [0,1] (sigmoid of whitened values by default) to solve the optimal transport problem;
cost_fn, used in
PointCloud, that defines the ground cost function to transport from
num_targetstarget values (squared Euclidean distance by default, see
pointcloud.pyfor more details);
epsilonvalues as well as other parameters to shape the
- Return type
A jnp.ndarray[batch | topk, dim].