code: def sinkhorn_matcher(
      cost: jnp.ndarray,
      rng: Optional[jnp.ndarray] = None,
      epsilon: float = 0.001,
      init: float = 50,
      decay: float = 0.9,
      num_iters: int = 1000,
      num_permutations: int = 100,
      threshold: float = 1e-2,
      chg_momentum_from: int = 100,
  ):
      """Computes Sinkhorn Matching on cost matrix for a batch of datapoints.
  
      Args:
        cost: Cost matrix for the matching of shape [B, N, N].
        rng: Random generator for sampling.
        epsilon: Level of entropic regularization wanted.
        init: Multiplier for epsilon decay at the first iteration.
        decay: How much to decay epsilon between two iterations.
        num_iters: Number of Sinkhorn iterations.
        num_permutations: Number of random permutations to sample for
          selecting the best.
        threshold: Convergence threshold for Sinkhorn algorithm.
        chg_momentum_from: Iteration from which to trigger the momemtum in Sinkhorn.
  
      Returns:
        An assignment of size [B, 2, N].
      """
  
      def coupling_fn(c):
          geom = geometry.Geometry(
              cost_matrix=c, epsilon=epsilon, init=init, decay=decay
          )
          return transport.solve(
              geom,
              max_iterations=num_iters,
              chg_momentum_from=chg_momentum_from,
              threshold=threshold,
          ).matrix
  
      coupling = jax.vmap(coupling_fn)(cost)
  
      if rng is None:
          # Use fixed key to make sampling deterministic.
          rng = jax.random.PRNGKey(0)
  
      permutation = sample_best_permutation(
          rng, coupling, cost, num_permutations
      )
      permutation = jnp.array(permutation, dtype=jnp.int32)
      grid = np.stack(
          np.meshgrid(
              np.arange(cost.shape[1], dtype=np.int32),
              np.arange(cost.shape[2], dtype=np.int32),
              indexing="ij",
          ),
          axis=0,
      )
      indices = jnp.einsum("bnm,2nm->b2n", permutation, grid)
  
      return indices
  
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 4 for calling a global function with example usage: sinkhorn_matcher(cost, rng, epsilon, init, decay, num_iters, num_permutations, threshold, chg_momentum_from) and returns: indices, transport