def reorder_histogram(bucket_vocab: tf.Tensor, counts: tf.Tensor,
                      boundary_size: int) -> tf.Tensor:
  """Return the histogram counts in indexed order, and zero out missing values.

  The count_elements analyzer returns counts in alphanumeric order, only for the
  values that are present. To construct a well-formed histogram, we need to
  rearrange them in numerical order, and fill in the missing values.

  Ex: The data contains values in the following form: [0, 1, 0, 1, 0, 3, 0, 1]
  bucket_indices happen to be the same as these values, and
  count_elements(tf.strings.as_string(bucket_indices)) returns:
    bucket_vocab=['1', '3', '0'],
    counts=[3, 1, 4]

  If boundaries=[0, 1, 2, 3, 4], we expect counts=[4, 3, 0, 1, 0],
  which this function will return.

  Args:
    bucket_vocab: A `Tensor` that names the buckets corresponding to the count
      information returned.
    counts: A `Tensor` that matches the bucket_vocab.
    boundary_size: A scalar that provides information about how big the returned
      counts should be.

  Returns:
    counts: A `Tensor` of size boundary_size corresponding to counts of all
        available buckets.
  """
  if bucket_vocab.dtype == tf.string:
    bucket_vocab = tf.strings.to_number(bucket_vocab, tf.int32)
  # counts/bucket_vocab may be out of order and missing values (empty buckets).
  ordering = tf.argsort(
      tf.concat([bucket_vocab,
                 tf.sets.difference([tf.range(boundary_size)],
                                    [bucket_vocab]).values], axis=-1))
  counts = tf.pad(counts, [[0, boundary_size - tf.size(counts)]])
  return tf.gather(counts, ordering)
Return the histogram counts in indexed order, and zero out missing values.

  The count_elements analyzer returns counts in alphanumeric order, only for the
  values that are present. To construct a well-formed histogram, we need to
  rearrange them in numerical order, and fill in the missing values.

  Ex: The data contains values in the following form: [0, 1, 0, 1, 0, 3, 0, 1]
  bucket_indices happen to be the same as these values, and
  count_elements(tf.strings.as_string(bucket_indices)) returns:
    bucket_vocab=['1', '3', '0'],
    counts=[3, 1, 4]

  If boundaries=[0, 1, 2, 3, 4], we expect counts=[4, 3, 0, 1, 0],
  which this function will return.

  Args:
    bucket_vocab: A `Tensor` that names the buckets corresponding to the count
      information returned.
    counts: A `Tensor` that matches the bucket_vocab.
    boundary_size: A scalar that provides information about how big the returned
      counts should be.

  Returns:
    counts: A `Tensor` of size boundary_size corresponding to counts of all
        available buckets.
  
tensorflow_transform/tf_utils.py