def reduce_batch_weighted_counts(
    x: common_types.TensorType,
    weights: Optional[tf.Tensor] = None,
    force: bool = False,
    filter_regex: Optional[str] = None) -> ReducedBatchWeightedCounts:
  """Performs batch-wise reduction to produce (possibly weighted) counts.

  Args:
    x: Input `Tensor` or `CompositeTensor`.
    weights: (Optional) Input weights.
    force: If True, reduces input tensor without weights to unique elements and
      counts.
    filter_regex: (Optional) Regex that matches tokens that have to be filtered
      out. May only be specified if `x` has string dtype.

  Returns:
    a named tuple of...
      The unique values in x
      The sum of the weights for each unique value in x if weights are provided,
        else None
  """
  if isinstance(x, tf.SparseTensor):
    x = x.values
  elif isinstance(x, tf.RaggedTensor):
    x = x.flat_values
  flat_x = tf.reshape(x, [-1])
  filter_fn = _make_regex_filter_fn(flat_x, filter_regex)
  flat_x = filter_fn(flat_x)
  if weights is None:
    if force:
      unique, _, counts = tf.unique_with_counts(flat_x)
      return ReducedBatchWeightedCounts(unique, None, None, counts)
    else:
      # TODO(b/112916494): Always do batch wise reduction once possible.
      return ReducedBatchWeightedCounts(flat_x, None, None, None)
  weights = filter_fn(tf.reshape(weights, [-1]))
  unique_x_values, unique_idx, _ = tf.unique_with_counts(
      flat_x, out_idx=tf.int64)
  summed_weights_per_x = tf.math.unsorted_segment_sum(
      weights, unique_idx, tf.size(input=unique_x_values))
  return ReducedBatchWeightedCounts(unique_x_values, summed_weights_per_x, None,
                                    None)
Performs batch-wise reduction to produce (possibly weighted) counts.

  Args:
    x: Input `Tensor` or `CompositeTensor`.
    weights: (Optional) Input weights.
    force: If True, reduces input tensor without weights to unique elements and
      counts.
    filter_regex: (Optional) Regex that matches tokens that have to be filtered
      out. May only be specified if `x` has string dtype.

  Returns:
    a named tuple of...
      The unique values in x
      The sum of the weights for each unique value in x if weights are provided,
        else None
  
tensorflow_transform/tf_utils.py