|View source on GitHub|
Computes a matrix representing a continuous relaxation of sorting.
tfp.experimental.substrates.jax.math.soft_sorting_matrix( x, temperature, name=None )
Given a vector
x, there exists a permutation matrix
P_x, when applied to
x sorted in decreasing order. Here, we compute a continuous
P_x, parameterized by
temperature. This continuous
relaxation satisfies the property that it is a unimodal row-stochastic matrix,
meaning that all entries are non-negative, all rows sum to 1., and there is a
unique maximum entry in each column. The unique maximum entry will correspond
to the location of a
1 in the exact sorting permutation.
Complexity: Given a vector
x of size
N, this operation will take
This is also known as a Neural sort in .
Tensor. Argument to compute the relaxed sorting matrix with respect to. The relaxed permutation is computed with respect to the last axis.
. Whentemperature` approaches zero, this will retrieve the exact permutation matrix corresponding to sorting from largest to smallest.
strname prefixed to Ops created by this function. Default value:
soft_sort: A unimodal row-stochastic matrix. Applying this matrix on x will in the limit of low temperature, sort it.
: Aditya Grover, Eric Wang, Aaron Zweig, Stefano Ermon. Stochastic Optimization of Sorting Networks via Continuous Relaxations. https://arxiv.org/abs/1903.08850