def _swap_m_with_i(vecs, m, i):
vecs = tf.convert_to_tensor(value=vecs, dtype=tf.int64, name='<STR_LIT>')<EOL>m = tf.convert_to_tensor(value=m, dtype=tf.int64, name='<STR_LIT:m>')<EOL>i = tf.convert_to_tensor(value=i, dtype=tf.int64, name='<STR_LIT:i>')<EOL>trailing_elts = tf.broadcast_to(<EOL>tf.range(m + <NUM_LIT:1>,<EOL>prefer_static.shape(vecs, out_type=tf.int64)[-<NUM_LIT:1>]),<EOL>prefer_static.shape(vecs[..., m + <NUM_LIT:1>:]))<EOL>shp = prefer_static.shape(trailing_elts)<EOL>trailing_elts = tf.where(<EOL>tf.equal(trailing_elts, tf.broadcast_to(i, shp)),<EOL>tf.broadcast_to(tf.gather(vecs, [m], axis=-<NUM_LIT:1>), shp),<EOL>tf.broadcast_to(vecs[..., m + <NUM_LIT:1>:], shp))<EOL>vecs_shape = vecs.shape<EOL>vecs = tf.concat([<EOL>vecs[..., :m],<EOL>tf.gather(vecs, i, batch_dims=prefer_static.rank(vecs) - <NUM_LIT:1>), trailing_elts<EOL>], axis=-<NUM_LIT:1>)<EOL>tensorshape_util.set_shape(vecs, vecs_shape)<EOL>return vecs<EOL>
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.