Announcing the TensorFlow Dev Summit 2020

# tfp.experimental.substrates.jax.math.soft_sorting_matrix

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` 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 .

#### 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`. When`temperature` approaches zero, this will retrieve the exact permutation matrix corresponding to sorting from largest to smallest.
• `name`: Python `str` name 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.

: Aditya Grover, Eric Wang, Aaron Zweig, Stefano Ermon. Stochastic Optimization of Sorting Networks via Continuous Relaxations. https://arxiv.org/abs/1903.08850