code: def soft_sorting_matrix(x, temperature, name=None):
      """Computes a matrix representing a continuous relaxation of sorting.
  
      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`. 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.
  
      #### References
  
      [1]: Aditya Grover, Eric Wang, Aaron Zweig, Stefano Ermon.
           Stochastic Optimization of Sorting Networks via Continuous Relaxations.
           https://arxiv.org/abs/1903.08850
      """
      with tf.name_scope(name or "soft_sorting_matrix"):
          dtype = dtype_util.common_dtype(
              [temperature, x], dtype_hint=tf.float32
          )
          temperature = tf.convert_to_tensor(
              temperature, name="temperature", dtype=dtype
          )
          x = tf.convert_to_tensor(x, name="x", dtype=dtype)
          n = tf.shape(x)[-1]
          y = x[..., tf.newaxis]
          pairwise_distances = tf.abs(y - tf.linalg.matrix_transpose(y))
          scaling = tf.cast(
              tf.range(n - 1, -(n - 1) - 1, delta=-2), dtype=dtype
          )
          p_logits = tf.linalg.matrix_transpose(
              tf.matmul(y, scaling[tf.newaxis, ...])
              - tf.reduce_sum(pairwise_distances, axis=-1)[
                  ..., tf.newaxis
              ]
          )
          y = tf.nn.softmax(p_logits / temperature, axis=-1)
          return y
  
imports: import collections
  import functools
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import custom_gradient
  from tensorflow_probability.python.internal import distribute_lib
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow_probability.python.internal import variadic_reduce
  from tensorflow_probability.python.math.scan_associative import scan_associative
  from tensorflow.python.util import deprecation
label_desc: Use image node_id 21 for calling a global function with example usage: soft_sorting_matrix(x, temperature, name) and returns: y