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(value=vecs, dtype=tf.int64, name='vecs')
  m = tf.convert_to_tensor(value=m, dtype=tf.int64, name='m')
  i = tf.convert_to_tensor(value=i, dtype=tf.int64, name='i')
  trailing_elts = tf.broadcast_to(
      tf.range(m + 1,
               prefer_static.shape(vecs, out_type=tf.int64)[-1]),
      prefer_static.shape(vecs[..., m + 1:]))
  shp = prefer_static.shape(trailing_elts)
  trailing_elts = tf.where(
      tf.equal(trailing_elts, tf.broadcast_to(i, shp)),
      tf.broadcast_to(tf.gather(vecs, [m], axis=-1), shp),
      tf.broadcast_to(vecs[..., m + 1:], shp))
  # TODO(bjp): Could we use tensor_scatter_nd_update?
  vecs_shape = vecs.shape
  vecs = tf.concat([
      vecs[..., :m],
      tf.gather(vecs, i, batch_dims=prefer_static.rank(vecs) - 1), trailing_elts
  ], axis=-1)
  tensorshape_util.set_shape(vecs, vecs_shape)
  return vecs
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(value=vecs, dtype=tf.int64, name='vecs')
  m = tf.convert_to_tensor(value=m, dtype=tf.int64, name='m')
  i = tf.convert_to_tensor(value=i, dtype=tf.int64, name='i')
  trailing_elts = tf.broadcast_to(
      tf.range(m + 1,
               prefer_static.shape(vecs, out_type=tf.int64)[-1]),
      prefer_static.shape(vecs[..., m + 1:]))
  shp = prefer_static.shape(trailing_elts)
  trailing_elts = tf.where(
      tf.equal(trailing_elts, tf.broadcast_to(i, shp)),
      tf.broadcast_to(tf.gather(vecs, [m], axis=-1), shp),
      tf.broadcast_to(vecs[..., m + 1:], shp))
  # TODO(bjp): Could we use tensor_scatter_nd_update?
  vecs_shape = vecs.shape
  vecs = tf.concat([
      vecs[..., :m],
      tf.gather(vecs, i, batch_dims=prefer_static.rank(vecs) - 1), trailing_elts
  ], axis=-1)
  tensorshape_util.set_shape(vecs, vecs_shape)
  return vecs
def _swap_m_with_i ( vecs , m , i ) : vecs = tf . convert_to_tensor ( value = vecs , dtype = tf . int64 , name = 'vecs' ) m = tf . convert_to_tensor ( value = m , dtype = tf . int64 , name = 'm' ) i = tf . convert_to_tensor ( value = i , dtype = tf . int64 , name = 'i' ) trailing_elts = tf . broadcast_to ( tf . range ( m + 1 , prefer_static . shape ( vecs , out_type = tf . int64 ) [ - 1 ] ) , prefer_static . shape ( vecs [ ... , m + 1 : ] ) ) shp = prefer_static . shape ( trailing_elts ) trailing_elts = tf . where ( tf . equal ( trailing_elts , tf . broadcast_to ( i , shp ) ) , tf . broadcast_to ( tf . gather ( vecs , [ m ] , axis = - 1 ) , shp ) , tf . broadcast_to ( vecs [ ... , m + 1 : ] , shp ) ) # TODO(bjp): Could we use tensor_scatter_nd_update? vecs_shape = vecs . shape vecs = tf . concat ( [ vecs [ ... , : m ] , tf . gather ( vecs , i , batch_dims = prefer_static . rank ( vecs ) - 1 ) , trailing_elts ] , axis = - 1 ) tensorshape_util . set_shape ( vecs , vecs_shape ) return vecs
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.
Swaps m and i on axis - 1 . ( Helper for pivoted_cholesky . )