code: def _sample_from_edpp(eigenvectors, vector_onehot, seed):
      """Samples a batch of subsets from a DPP given pre-selected elementary DPPs.
  
      Recall that an elementary DPP is a DPP with eigenvalues all exactly 0 or 1.
      This function implements the second step of standard sampling algorithm for
      DPPs, by sampling subsets based on the E-DPPs obtained by selecting
      `vector_onehot` against the DPP's original eigenvectors.
  
      Args:
        eigenvectors: A Tensor of `float32` of shape `[..., num_points, num_vecs]`
          representing the eigenvectors of a DPP's L-ensemble matrix, eigenvectors
          in columns. Generally, `num_vecs == num_points`; we name separately to
          distinguish axes.
        vector_onehot:  A Tensor of shape `[..., n_vecs]` whose innermost
          dimension corresponds to 1-hot subset encodings. The subsets represent the
          subset of eigenvectors of the original DPP that define an elementary DPP.
        seed: PRNG seed; see `tfp.random.sanitize_seed` for details.
  
      Returns:
        samples: A many-hot `bool` Tensor of shape `[..., n_points]`
          representing a batch of 1-hot subset encodings.
      """
      with tf.name_scope("sample_from_edpp"):
          seed = samplers.sanitize_seed(seed)
          # Sort the 1's to the front, and sort corresponding eigenvectors, then mask.
          vector_onehot = tf.cast(vector_onehot, eigenvectors.dtype)
          vector_indices = tf.argsort(
              vector_onehot, axis=-1, direction="DESCENDING"
          )
          vector_onehot = tf.gather(
              vector_onehot,
              vector_indices,
              axis=-1,
              batch_dims=len(vector_indices.shape) - 1,
          )
          eigenvectors = tf.gather(
              eigenvectors,
              vector_indices,
              axis=-1,
              batch_dims=len(vector_indices.shape) - 1,
          )
          eigenvectors = (
              eigenvectors * vector_onehot[..., tf.newaxis, :]
          )
          sample_size = tf.reduce_sum(
              tf.cast(vector_onehot, tf.int32), axis=-1
          )
          max_sample_size = tf.reduce_max(sample_size)
  
          d = ps.shape(eigenvectors)[-2]
          n = ps.shape(eigenvectors)[-1]
  
          # Slice eigvecs to do less work in eager/non-XLA modes.
          if (
              FAST_PATH_ENABLED
              and not JAX_MODE
              and (
                  tf.executing_eagerly()
                  or not control_flow_util.GraphOrParentsInXlaContext(
                      tf1.get_default_graph()
                  )
              )
          ):
              # We can save some work in non-XLA contexts by reducing the size of the
              # eigenvectors.
              eigenvectors = eigenvectors[..., :max_sample_size]
              n = max_sample_size
  
          def cond(i, *_):
              return i < max_sample_size
  
          def body(i, vecs, cur_sample, seed):
              sample_seed, next_seed = samplers.split_seed(seed)
              # squared norm at each coord across active subspace
              is_active = i < sample_size
              coord_prob = tf.reduce_sum(tf.square(vecs), axis=-1)
              coord_logits = tf.where(
                  is_active[..., tf.newaxis],
                  tf.math.log(coord_prob),
                  0.0,
              )
  
              idx = categorical.Categorical(logits=coord_logits).sample(
                  seed=sample_seed
              )
              new_vecs = tf.where(
                  (
                      tf.range(n)
                      < sample_size[..., tf.newaxis, tf.newaxis] - i - 1
                  )
                  & ~cur_sample[..., tf.newaxis],
                  _orthogonal_complement_e_i(
                      vecs,
                      i=tf.where(is_active, idx, 0),
                      gram_schmidt_iters=max_sample_size - i,
                  ),
                  0.0,
              )
              # Since range(n) may have unknown shape in the stmt above, we clarify.
              tensorshape_util.set_shape(new_vecs, vecs.shape)
              vecs = tf.where(
                  is_active[..., tf.newaxis, tf.newaxis], new_vecs, vecs
              )
              cur_sample = cur_sample | (
                  tf.equal(tf.range(d), idx[..., tf.newaxis])
                  & is_active[..., tf.newaxis]
              )
              return i + 1, vecs, cur_sample, next_seed
  
          _, _, sample, _ = tf.while_loop(
              cond,
              body,
              (
                  tf.zeros([], tf.int32, name="i"),
                  eigenvectors,
                  tf.zeros(ps.shape(eigenvectors)[:-1], dtype=tf.bool),
                  seed,
              ),
          )
  
          return tf.cast(sample, tf.int32)
  
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 3 for calling a global function with example usage: _sample_from_edpp(eigenvectors, vector_onehot, seed) and returns: tf, unknown, unknown,  vecs,  cur_sample,  next_seed