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
Compute the element bucket and update the histogram.
language/boolq/utils/ops.py