code: def idx2permutation(row_ind, col_ind):
      """Constructs a permutation matrix from the column and row indices of ones."""
      bs, dim = row_ind.shape[:2]
      perm = jnp.zeros(shape=(bs, dim, dim), dtype="float32")
      perm = jax.vmap(
          (lambda x, idx, y: x.at[idx].set(y)), (0, 0, None)
      )(perm, (row_ind, col_ind), 1.0)
      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 1 for calling a global function with example usage: idx2permutation(row_ind, col_ind) and returns: perm