def per_voxel_point_sample_segment_func(data, segment_ids, num_segments,
                                        num_samples_per_voxel):
  """Samples features from the points within each voxel.

  Args:
    data: A tf.float32 tensor of size [N, F].
    segment_ids: A tf.int32 tensor of size [N].
    num_segments: Number of segments.
    num_samples_per_voxel: Number of features to sample per voxel. If the voxel
      has less number of points in it, the point features will be padded by 0.

  Returns:
    A tf.float32 tensor of size [num_segments, num_samples_per_voxel, F].
    A tf.int32 indices of size [N, num_samples_per_voxel].
  """
  num_channels = data.get_shape().as_list()[1]
  if num_channels is None:
    raise ValueError('num_channels is None.')
  n = tf.shape(segment_ids)[0]

  def _body_fn(i, indices_range, indices):
    """Computes the indices of the i-th point feature in each segment."""
    indices_i = tf.math.unsorted_segment_max(
        data=indices_range, segment_ids=segment_ids, num_segments=num_segments)
    indices_i_positive_mask = tf.greater(indices_i, 0)
    indices_i_positive = tf.boolean_mask(indices_i, indices_i_positive_mask)
    boolean_mask = tf.scatter_nd(
        indices=tf.cast(
            tf.expand_dims(indices_i_positive - 1, axis=1), dtype=tf.int64),
        updates=tf.ones_like(indices_i_positive, dtype=tf.int32),
        shape=(n,))
    indices_range *= (1 - boolean_mask)
    indices_i *= tf.cast(indices_i_positive_mask, dtype=tf.int32)
    indices_i = tf.pad(
        tf.expand_dims(indices_i, axis=1),
        paddings=[[0, 0], [i, num_samples_per_voxel - i - 1]])
    indices += indices_i
    i = i + 1
    return i, indices_range, indices

  cond = lambda i, indices_range, indices: i < num_samples_per_voxel

  (_, _, indices) = tf.while_loop(
      cond=cond,
      body=_body_fn,
      loop_vars=(tf.constant(0, dtype=tf.int32), tf.range(n) + 1,
                 tf.zeros([num_segments, num_samples_per_voxel],
                          dtype=tf.int32)))

  data = tf.pad(data, paddings=[[1, 0], [0, 0]])
  voxel_features = tf.gather(data, tf.reshape(indices, [-1]))
  return tf.reshape(voxel_features,
                    [num_segments, num_samples_per_voxel, num_channels])