def bucket_by_quantiles(len_fn, batch_size, n_buckets, hist_bounds):
  """Dynamically bucket a `tf.data.Dataset` based on the element's length.

  Keeps track of a histogram of the input elements lengths, and yields batches
  of examples that belong to the same length quantile.

  Useful in cases where you want to bucket data, but don't know what the
  optimal bucketing ranges should be

  Args:
    len_fn: Function mapping elements in the dataset to an integer length
    batch_size: Maximum size of the output batches
    n_buckets: Number of quantiles to break the data into
    hist_bounds: List of integer bounds to use when building the histograms,
      should cover a range so that at most a single quantile of elements are
      lower/higher then the bucket range. More fine-grained buckets will make
      the histogram more precise, but adds to the computational overhead

  Raises:
    ValueError: If `hist_bounds` or `len_fn` are invalid
  Returns:
    A function that can be used with tf.data.Dataset.apply to batch a dataset
  """
  n_hist_binds = len(hist_bounds)

  if n_hist_binds < n_buckets:
    raise ValueError("Requested %d buckets, but only have %d histogram bins" %
                     (n_buckets, n_hist_binds))
  if any(hist_bounds[i] >= hist_bounds[i + 1] for i in range(n_hist_binds - 1)):
    raise ValueError("Bins must be descending")

  # Need to use `use_resource = True` to make this work correctly
  # within tf.data.Dataset
  hist_counts = tf.get_local_variable(
      "hist-counts",
      n_hist_binds + 1,
      tf.int64,
      tf.zeros_initializer(),
      use_resource=True)
  hist_bounds = tf.constant(hist_bounds, tf.int64)

  def bucket_fn(x):
    """Compute the element bucket and update the histogram."""
    ix = len_fn(x)
    if ix.dtype == tf.int32:
      ix = tf.to_int64(ix)
    elif ix.dtype != tf.int64:
      raise ValueError("Len function returned a non-int")

    adds_to_bins = tf.to_int64(tf.greater(hist_bounds, ix))
    # pad with a 1 for the "larger than all" bin
    adds_to_bins = tf.pad(adds_to_bins, [[0, 1]], constant_values=1)
    new_counts = tf.assign_add(hist_counts, adds_to_bins)
    bin_ix = n_hist_binds - tf.reduce_sum(adds_to_bins)
    # Computes the quantile based on the counts of the exammple's bucket
    bucket_ix = tf.floordiv(((n_buckets - 1) * new_counts[bin_ix]),
                            new_counts[-1])
    return bucket_ix

  def reduce_fn(_, x):
    return x.padded_batch(batch_size, dataset_ops.get_legacy_output_shapes(x))

  return tf.data.experimental.group_by_window(bucket_fn, reduce_fn, batch_size)
Dynamically bucket a `tf.data.Dataset` based on the element's length.

  Keeps track of a histogram of the input elements lengths, and yields batches
  of examples that belong to the same length quantile.

  Useful in cases where you want to bucket data, but don't know what the
  optimal bucketing ranges should be

  Args:
    len_fn: Function mapping elements in the dataset to an integer length
    batch_size: Maximum size of the output batches
    n_buckets: Number of quantiles to break the data into
    hist_bounds: List of integer bounds to use when building the histograms,
      should cover a range so that at most a single quantile of elements are
      lower/higher then the bucket range. More fine-grained buckets will make
      the histogram more precise, but adds to the computational overhead

  Raises:
    ValueError: If `hist_bounds` or `len_fn` are invalid
  Returns:
    A function that can be used with tf.data.Dataset.apply to batch a dataset
  
language/boolq/utils/ops.py