def pointcloud_to_sparse_voxel_grid(points, features, num_valid_points,
                                    grid_cell_size, voxels_pad_or_clip_size,
                                    segment_func):
  """Converts a pointcloud into a voxel grid.

  This function calls the `pointcloud_to_sparse_voxel_grid_unbatched`
  function above in a while loop to map a batch of points to a batch of voxels.

  Args:
    points: A tf.float32 tensor of size [batch_size, N, 3].
    features: A tf.float32 tensor of size [batch_size, N, F].
    num_valid_points: A tf.int32 tensor of size [num_batches] containing the
      number of valid points in each batch example.
    grid_cell_size: A tf.float32 tensor of size [3].
    voxels_pad_or_clip_size: Number of target voxels to pad or clip to. If None,
      it will not perform the padding.
    segment_func: A tensorflow function that operates on segments. Examples are
      one of tf.math.unsorted_segment_{min/max/mean/prod/sum}.

  Returns:
    voxel_features: A tf.float32 tensor of size [batch_size, N', F]
      or [batch_size, N', G, F] where G is the number of points sampled per
      voxel.
    voxel_indices: A tf.int32 tensor of size [batch_size, N', 3].
    num_valid_voxels: A tf.int32 tensor of size [batch_size].
    segment_ids: A size [batch_size, N] tf.int32 tensor of IDs for each point
      indicating which (flattened) voxel cell its data was mapped to.
    voxel_start_location: A size [batch_size, 3] tf.float32 tensor of voxel
      start locations.

  Raises:
    ValueError: If pooling method is unknown.
  """
  batch_size = points.get_shape().as_list()[0]
  if batch_size is None:
    batch_size = tf.shape(points)[0]
  num_points = tf.shape(points)[1]

  def fn(i):
    """Map function."""
    num_valid_points_i = num_valid_points[i]
    points_i = points[i, :num_valid_points_i, :]
    features_i = features[i, :num_valid_points_i, :]
    voxel_features_i, voxel_indices_i, segment_ids_i, voxel_start_location_i = (
        pointcloud_to_sparse_voxel_grid_unbatched(
            points=points_i,
            features=features_i,
            grid_cell_size=grid_cell_size,
            segment_func=segment_func))
    num_valid_voxels_i = tf.shape(voxel_features_i)[0]
    (voxel_features_i, voxel_indices_i, num_valid_voxels_i,
     segment_ids_i) = _pad_or_clip_voxels(
         voxel_features=voxel_features_i,
         voxel_indices=voxel_indices_i,
         num_valid_voxels=num_valid_voxels_i,
         segment_ids=segment_ids_i,
         voxels_pad_or_clip_size=voxels_pad_or_clip_size)
    segment_ids_i = tf.pad(
        segment_ids_i, paddings=[[0, num_points - num_valid_points_i]])
    return (voxel_features_i, voxel_indices_i, num_valid_voxels_i,
            segment_ids_i, voxel_start_location_i)

  return tf.map_fn(
      fn=fn,
      elems=tf.range(batch_size),
      dtype=(tf.float32, tf.int32, tf.int32, tf.int32, tf.float32))