Batching function.

def batch_fn(dataset, training, shapes, target_names,
             batch_size=32, eval_batch_size=32, bucket_batch_length=32,
             bucket_max_length=256, bucket_min_length=8,
             bucket_length_step=1.1, buckets=None):
  """Batching function."""
  del target_names
  # If bucketing is not specified, check if target shapes are variable.
  cur_batch_size = batch_size if training else eval_batch_size
  if buckets is None:
    variable_target_shapes = False
    target_shape = shapes[1]
    for dim in target_shape:
      if dim is None:
        variable_target_shapes = True
    tf.logging.info("Heuristically setting bucketing to %s based on shapes "
                    "of target tensors." % variable_target_shapes)
    if variable_target_shapes:
      batch_size_per_token = cur_batch_size * bucket_batch_length
      scheme = data_reader.batching_scheme(batch_size_per_token,
                                           bucket_max_length,
                                           bucket_min_length,
                                           bucket_length_step,
                                           drop_long_sequences=training)
      buckets = (scheme["boundaries"], scheme["batch_sizes"])

  if buckets:
    tf.logging.info("Bucketing with buckets %s." % str(buckets))
    def example_length(_, target):
      return tf.shape(target)[0]
    boundaries, batch_sizes = buckets
    dataset = dataset.apply(tf.data.experimental.bucket_by_sequence_length(
        example_length, boundaries, batch_sizes))
  else:
    dataset = dataset.padded_batch(cur_batch_size, shapes)
  return dataset