code: def scheduled_sample_count(ground_truth_x,
                             generated_x,
                             batch_size,
                             scheduled_sample_var):
    
    num_ground_truth = scheduled_sample_var
    idx = tf.random_shuffle(tf.range(batch_size))
    ground_truth_idx = tf.gather(idx, tf.range(num_ground_truth))
    generated_idx = tf.gather(idx, tf.range(num_ground_truth, batch_size))
  
    ground_truth_examps = tf.gather(ground_truth_x, ground_truth_idx)
    generated_examps = tf.gather(generated_x, generated_idx)
  
    output = tf.dynamic_stitch([ground_truth_idx, generated_idx],
                               [ground_truth_examps, generated_examps])
    # if batch size is known set it.
    if isinstance(batch_size, int):
      output.set_shape([batch_size] + common_layers.shape_list(output)[1:])
    return output
docs: Sample batch with specified mix of groundtruth and generated data points.
  
    Args:
      ground_truth_x: tensor of ground-truth data points.
      generated_x: tensor of generated data points.
      batch_size: batch size
      scheduled_sample_var: number of ground-truth examples to include in batch.
    Returns:
      New batch with num_ground_truth sampled from ground_truth_x and the rest
      from generated_x.