def voxels_to_points(voxels, segment_ids):
  """Convert voxels back to points given their segment id.

  Args:
    voxels: A tf.float32 tensor representing a voxel grid. Expect shape
      [x, y, z, f].
    segment_ids: A tf.int32 tensor representing the segment id of each point
      in the original pointcloud we want to project voxel features back to.
  Returns:
    point_features: A tf.float32 tensor of shape [N, f] where each point
      now has the features in the associated voxel cell.
  """
  flattened_voxels = tf.reshape(voxels, shape=(-1, voxels.shape[-1]))
  return tf.gather(flattened_voxels, segment_ids)