code: def _swap_m_with_i(vecs, m, i):
      """Swaps `m` and `i` on axis -1. (Helper for pivoted_cholesky.)
  
      Given a batch of int64 vectors `vecs`, scalar index `m`, and compatibly shaped
      per-vector indices `i`, this function swaps elements `m` and `i` in each
      vector. For the use-case below, these are permutation vectors.
  
      Args:
        vecs: Vectors on which we perform the swap, int64 `Tensor`.
        m: Scalar int64 `Tensor`, the index into which the `i`th element is going.
        i: Batch int64 `Tensor`, shaped like vecs.shape[:-1] + [1]; the index into
          which the `m`th element is going.
  
      Returns:
        vecs: The updated vectors.
      """
      vecs = tf.convert_to_tensor(vecs, dtype=tf.int64, name="vecs")
      m = tf.convert_to_tensor(m, dtype=tf.int64, name="m")
      i = tf.convert_to_tensor(i, dtype=tf.int64, name="i")
      trailing_elts = tf.broadcast_to(
          tf.range(m + 1, ps.shape(vecs, out_type=tf.int64)[-1]),
          ps.shape(vecs[..., m + 1 :]),
      )
      trailing_elts = tf.where(
          tf.equal(trailing_elts, i),
          tf.gather(vecs, [m], axis=-1),
          vecs[..., m + 1 :],
      )
      # TODO(bjp): Could we use tensor_scatter_nd_update?
      vecs_shape = vecs.shape
      vecs_rank = tensorshape_util.rank(vecs_shape)
      if vecs_rank is None:
          raise NotImplementedError(
              "Input vector to swap must have statically known rank. If you cannot "
              "provide static rank please contact core TensorFlow team and request "
              "that `tf.gather` `batch_dims` argument support `tf.Tensor`-valued "
              "inputs."
          )
      vecs = tf.concat(
          [
              vecs[..., :m],
              tf.gather(vecs, i, batch_dims=int(vecs_rank) - 1),
              trailing_elts,
          ],
          axis=-1,
      )
      tensorshape_util.set_shape(vecs, vecs_shape)
      return vecs
  
imports: import numpy
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import assert_util
  from tensorflow_probability.python.internal import custom_gradient
  from tensorflow_probability.python.internal import distribution_util
  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.math import generic
label_desc: Use image node_id 3 for calling a global function with example usage: _swap_m_with_i(vecs, m, i) and returns: vecs