code: def sample_best_permutation(key, coupling, cost, num_trials=10):
      """Samples permutation matrices and returns the one with lowest cost.
  
      See **Convex Relaxations for Permutation Problems** paper for rough
      explanation of the algorithm.
  
      Args:
        key: jnp.ndarray; functions as a PRNG key.
        coupling: jnp.ndarray; has shape [N, N].
        cost: jnp.ndarray; has shape [N, N].
        num_trials: int; determines the amount of times we sample a permutation.
  
      Returns:
        Permutation matrix: jnp.ndarray of shape [N, N] of floating point type.
          This is the permutation matrix with lowest optimal transport cost.
      """
      vec_sample_permutation = jax.vmap(
          sample_permutation, in_axes=(0, None), out_axes=0
      )
      key = jax.random.split(key, num_trials)
      perms = vec_sample_permutation(key, coupling)
  
      # Pick the permutation with minimal ot cost
      ot = jnp.einsum("nbij,bij->nb", perms, cost)
      min_idx = jnp.argmin(ot, axis=0)
      out_perm = jax.vmap(jnp.take, (1, 0, None))(perms, min_idx, 0)
      return out_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 3 for calling a global function with example usage: sample_best_permutation(key, coupling, cost, num_trials) and returns: out_perm