def extend_reduced_batch_with_y_counts(reduced_batch, y, weights=None):
  """Extend the ReducedBatchWeightedCounts with global counts for y.

  This is used to maintain an accurate count of global frequencies of each value
  in y. When x is multivalent, the sum over the summed_positive_per_x_and_y
  will over-count the occurrence of y. To keep track of the true distribution
  of y values, we add a sentinel value that tracks the global counts of each
  distinct value in y. This is useful for computing the mutual information
  between values in x and y.

  Args:
    reduced_batch: A ReducedBatchWeightedCounts instance.
    y: A `Tensor` representing a batch of y values.
    weights: Optional `Tensor` representing a batch of weight values.

  Returns:
    A new ReducedBatchWeightedCounts instance with sentinel values appended.
  """
  if isinstance(y, tf.SparseTensor):
    y_shape = y.dense_shape
    # Need to slice instead of directly using [0] as this runs in graph mode.
    shape = tf.slice(y_shape, [0], [1])
  else:
    shape = tf.shape(y)
  # Create a dummy sentinel token that is present in every record.
  if reduced_batch.unique_x.dtype.is_integer:
    sentinel_values = tf.cast(
        tf.fill(shape, GLOBAL_Y_COUNT_SENTINEL_INT), tf.int64
    )
  else:
    sentinel_values = tf.fill(shape, GLOBAL_Y_COUNT_SENTINEL_STRING)
  # Computing the batch reduction over this sentinel token will reduce to a
  # single sentinel value in sentinel_batch.unique_x, with the
  # summed_positive_per_x_and_y thus capturing the total summed positive per
  # value in y.
  sentinel_batch = reduce_batch_weighted_cooccurrences(
      sentinel_values, y, weights, extend_with_sentinel_counts=False)

  # Concatenate the sentinel counts with the existing reduced batch.
  return ReducedBatchWeightedCounts(
      unique_x=tf.concat([reduced_batch.unique_x, sentinel_batch.unique_x],
                         axis=0),
      summed_weights_per_x=tf.concat([
          reduced_batch.summed_weights_per_x,
          sentinel_batch.summed_weights_per_x
      ],
                                     axis=0),
      summed_positive_per_x_and_y=tf.concat([
          reduced_batch.summed_positive_per_x_and_y,
          sentinel_batch.summed_positive_per_x_and_y
      ],
                                            axis=0),
      counts_per_x=tf.concat(
          [reduced_batch.counts_per_x, sentinel_batch.counts_per_x], axis=0))
Extend the ReducedBatchWeightedCounts with global counts for y.

  This is used to maintain an accurate count of global frequencies of each value
  in y. When x is multivalent, the sum over the summed_positive_per_x_and_y
  will over-count the occurrence of y. To keep track of the true distribution
  of y values, we add a sentinel value that tracks the global counts of each
  distinct value in y. This is useful for computing the mutual information
  between values in x and y.

  Args:
    reduced_batch: A ReducedBatchWeightedCounts instance.
    y: A `Tensor` representing a batch of y values.
    weights: Optional `Tensor` representing a batch of weight values.

  Returns:
    A new ReducedBatchWeightedCounts instance with sentinel values appended.
  
tensorflow_transform/tf_utils.py