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].
Args | |
|---|---|
x
|
float Tensor. Argument to compute the relaxed sorting matrix with
respect to. The relaxed permutation is computed with respect to the last
axis.
|
temperature
|
Positive float Tensor. Whentemperatureapproaches zero,
this will retrieve the exact permutation matrix corresponding to sorting
from largest to smallest.
</td>
</tr><tr>
<td>name</td>
<td>
Pythonstrname prefixed to Ops created by this function.
Default value:None(i.e.,'soft_sorting_matrix'`).
|
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