code: def _orthogonal_complement_e_i(vectors, i, gram_schmidt_iters):
      """Computes a basis for the orthogonal complement to `e_i` in `span(vectors)`.
  
      The orthogonal complement of the coordinate vector `e_i` of the vector space
      `V` is the set of all vectors in `V` that are orthogonal to `e_i`.
  
      We compute this by first choosing a column `j` of `vectors` with non-zero in
      coordinate `i`. This vector (`col_j`) is subtracted from all other vectors
      with an appropriate weight to zero out row `i`. Finally, we orthonormalize
      using (modified) Gram-Schmidt. For performance reasons, the calling code
      specifies the G-S iteration count.
  
      For example, suppose we start with the matrix of column vectors:
  
      ```none
      [ 2  4  7 ]
      [ 4  2  4 ]
      [ 6  6  3 ]
      ```
  
      If we suppose `i = 1`, we are being asked to zero-out the middle row, i.e.
      orthogonalize with respect to the coordinate vector `e_1 = [0, 1, 0]^T`. We
      can do so by picking `j = argmax(mat[i, :])`, so `j = 0` in this case. Then,
      compute the appropriate weights that would zero out the row, i.e.
      `w=[1, 0.5, 1]` and subtract `mat[:, j:j+1] * w = [2, 4, 6]^T * [1, .5, 1]`.
      This yields the intermediate:
  
      ```none
      [ 2  4  7 ]   [ 2  1  2 ]   [ 0  3  5 ]
      [ 4  2  4 ] - [ 4  2  4 ] = [ 0  0  0 ]
      [ 6  6  3 ]   [ 6  3  6 ]   [ 0  3 -3 ]
      ```
  
      We rotate the zero column to the end, and finally return the result of
      applying Gram-Schmidt orthogonalization, i.e.
  
      ```none
      [ sqrt(.5)  sqrt(.5) 0 ]
      [     0        0     0 ]
      [ sqrt(.5) -sqrt(.5) 0 ]
      ```
  
      Args:
        vectors: A Tensor of vectors of shape `[..., d, n]` we are orthogonalizing.
        i: The coordinate (against dimension `d`) w.r.t. which we orthogonalize.
        gram_schmidt_iters: Number of iterations of Gram-Schmidt orthonormalization
          to run, generally `n_vectors - iter_num`. Since each iteration of sampling
          reduces the number of nonzero columns by one (in the `n` dim), this allows
          us to save iterations of orthonormalization work.
  
      Returns:
        orthogonal: A Tensor of shape `[..., d, n]` representing the subspace
          spanned by `vectors` that is orthogonal to `e_i`, the `i`-th coordinate
          vector. The tensor is orthonormalized. It contains at least one more zero
          row (`i`) and zero column than the input vectors (exactly one more if all
          nonzero columns of `vectors` are linearly independent).
      """
      i = tf.convert_to_tensor(i, dtype_hint=tf.int32)
      row_i = tf.gather(vectors, i, axis=-2, batch_dims=len(i.shape))
      j = tf.argmax(
          tf.abs(row_i), axis=-1
      )  # Max for numerical stability.
      col_j = tf.gather(vectors, j, axis=-1, batch_dims=len(j.shape))
      val_i_j = tf.gather(row_i, j, axis=-1, batch_dims=len(j.shape))
      weights = row_i / val_i_j[..., tf.newaxis]
      delta = weights[..., tf.newaxis, :] * col_j[..., :, tf.newaxis]
      result = vectors - delta
      # Rotate the new zero column to the end.
      d = ps.shape(vectors)[-2]
      n = ps.shape(vectors)[-1]
      mask_d = tf.not_equal(
          tf.range(d, dtype=i.dtype), i[..., tf.newaxis]
      )[..., tf.newaxis]
      shift_indices = tf.range(n, dtype=j.dtype)
      shift_indices = shift_indices + tf.cast(
          shift_indices >= j[..., tf.newaxis], j.dtype
      )
      shift_indices = tf.where(
          shift_indices >= tf.cast(n, j.dtype),
          j[..., tf.newaxis],
          shift_indices,
      )
      result = tf.gather(
          result,
          shift_indices,
          axis=-1,
          batch_dims=len(shift_indices.shape) - 1,
      )
      mask_n = tf.not_equal(tf.range(n), n - 1)
      result = tf.where(
          mask_d & mask_n, result, 0
      )  # Make exactly zero.
      # Orthonormalize. This is equivalent, but faster than tf.linalg.qr(result).q
      return gram_schmidt(result, gram_schmidt_iters)
  
imports: import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import softplus
  from tensorflow_probability.python.distributions import bernoulli
  from tensorflow_probability.python.distributions import categorical
  from tensorflow_probability.python.distributions import distribution
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import parameter_properties
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import reparameterization
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal import tensor_util
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow_probability.python.math.gram_schmidt import gram_schmidt
  from tensorflow.python.ops import control_flow_util
label_desc: Use image node_id 1 for calling a global function with example usage: _orthogonal_complement_e_i(vectors, i, gram_schmidt_iters) and returns: gram_schmidt