code: def sample_permutation(key, coupling):
      """Samples a permutation matrix from a doubly stochastic coupling matrix.
  
      CAREFUL: the couplings that come out of the Sinkhorn solver are not doubly
      stochastic but 1/dim * doubly_stochastic.
  
      See **Convex Relaxations for Permutation Problems** paper for rough
      explanation of the algorithm.
  
      Best to use by drawing multiple samples and picking the permutation with
      lowest cost as sometimes permutations seem to be drawn with high cost. The
      sample_best_permutation method does this.
  
      Args:
        key: jnp.ndarray; Functions as a PRNG key.
        coupling: jnp.ndarray; Has shape [N, N] which must have marginals such that
          coupling.sum(0) == 1. and coupling.sum(1) == 1. Note that in Sinkhorn we
          usually output couplings with marginals that sum to 1/N.
  
      Returns:
        Permutation matrix: jnp.ndarray of shape [N, N] of floating dtype.
      """
      bs, dim = coupling.shape[:2]
  
      # Random monotonic vector v without duplicates.
      v = jax.random.choice(
          key, 10 * dim, shape=(bs, dim), replace=False
      )
      v = jnp.sort(v, axis=-1) * 10.0
  
      w = jnp.einsum("bnm,bm->bn", coupling, v)
      # Sorting w will give the row indices of the permutation matrix.
      row_ind = jnp.argsort(w, axis=-1)
      col_ind = jnp.tile(jnp.arange(0, dim)[None, :], [bs, 1])
  
      # Compute permutation matrix from row and column indices.
      perm = idx2permutation(row_ind, col_ind)
      return perm
  
file_path: scenic/scenic/model_lib/matchers/sinkhorn.py
imports: from typing import Optional
  import jax
  import jax.numpy
  import numpy
  from ott.geometry import geometry
  from ott.tools import transport
label_desc: Use image node_id 2 for calling a global function with example usage: sample_permutation(key, coupling) and returns: perm