docs: def _batch_examples(dataset, batch_size, max_length):
    """Group examples by similar lengths, and return batched dataset.
  
    Each batch of similar-length examples are padded to the same length, and may
    have different number of elements in each batch, such that:
      group_batch_size * padded_length <= batch_size.
  
    This decreases the number of padding tokens per batch, which improves the
    training speed.
  
    Args:
      dataset: Dataset of unbatched examples.
      batch_size: Max number of tokens per batch of examples.
      max_length: Max number of tokens in an example input or target sequence.
  
    Returns:
      Dataset of batched examples with similar lengths.
    """
    # Get min and max boundary lists for each example. These are used to calculate
    # the `bucket_id`, which is the index at which:
    # buckets_min[bucket_id] <= len(example) < buckets_max[bucket_id]
    # Note that using both min and max lists improves the performance.
    buckets_min, buckets_max = _create_min_max_boundaries(max_length)
  
    # Create list of batch sizes for each bucket_id, so that
    # bucket_batch_size[bucket_id] * buckets_max[bucket_id] <= batch_size
    bucket_batch_sizes = [batch_size // x for x in buckets_max]
    # bucket_id will be a tensor, so convert this list to a tensor as well.
    bucket_batch_sizes = tf.constant(bucket_batch_sizes, dtype=tf.int64)
  
    def example_to_bucket_id(example_input, example_target):
      """Return int64 bucket id for this example, calculated based on length."""
      seq_length = _get_example_length((example_input, example_target))
  
      # TODO: investigate whether removing code branching improves performance.
      conditions_c = tf.logical_and(
          tf.less_equal(buckets_min, seq_length),
          tf.less(seq_length, buckets_max))
      bucket_id = tf.reduce_min(tf.where(conditions_c))
      return bucket_id
  
    def window_size_fn(bucket_id):
      """Return number of examples to be grouped when given a bucket id."""
      return bucket_batch_sizes[bucket_id]
  
    def batching_fn(bucket_id, grouped_dataset):
      """Batch and add padding to a dataset of elements with similar lengths."""
      bucket_batch_size = window_size_fn(bucket_id)
  
      # Batch the dataset and add padding so that all input sequences in the
      # examples have the same length, and all target sequences have the same
      # lengths as well. Resulting lengths of inputs and targets can differ.
      return grouped_dataset.padded_batch(bucket_batch_size, ([None], [None]))
  
    return dataset.apply(tf.contrib.data.group_by_window(
        key_func=example_to_bucket_id,
        reduce_func=batching_fn,
        window_size=None,
        window_size_func=window_size_fn))