View source on GitHub |
Computes a matrix representing a continuous relaxation of sorting.
tfp.math.soft_sorting_matrix(
x, temperature, name=None
)
Given a vector x
, there exists a permutation matrix P_x
, when applied to
x
gives x
sorted in decreasing order. Here, we compute a continuous
relaxation of 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 O(N**2)
time.
This is also known as a Neural sort in [1].
Returns | |
---|---|
soft_sort
|
A unimodal row-stochastic matrix. Applying this matrix on x will in the limit of low temperature, sort it. |
References
[1]: Aditya Grover, Eric Wang, Aaron Zweig, Stefano Ermon. Stochastic Optimization of Sorting Networks via Continuous Relaxations. https://arxiv.org/abs/1903.08850