def reduce_batch_weighted_cooccurrences(
    x_input: common_types.TensorType,
    y_input: Union[tf.Tensor, tf.SparseTensor],
    weights_input: Optional[tf.Tensor] = None,
    extend_with_sentinel_counts: bool = True,
    filter_regex: Optional[str] = None) -> ReducedBatchWeightedCounts:
  """Performs batch-wise reduction to produce weighted co-occurrences.

  Computes the weighted co-occurrence of each feature value in x, for each value
  in the range [0, max(y)). If extend_with_sentinel_counts is true, the return
  value will include an additional sentinel token (not in the true vocabulary)
  that is used to accumulate the global distribution of y values.

  Args:
    x_input: Input `Tensor` or `CompositeTensor`.
    y_input: Integer `Tensor` or `SparseTensor` (in case of
      multi-label/multi-task) with which to compute the co-occurrence with
      x_input.
    weights_input: (Optional) Weights input `Tensor`.
    extend_with_sentinel_counts: If True, the reduced batch will be extended
      a sentinel value that accumlate the total distribution of y values. Should
      be True except when called recursively with the sentinel value as input.
    filter_regex: (Optional) Regex that matches tokens that have to be filtered
      out. Can only be specified if `x_input` has string dtype.

  Returns:
    a namedtuple of...
    unique_x_values: the unique values in x
    summed_weights_per_x: sum of the weights for each unique value in x
    summed_positive_per_x_and_y: If tensor y is provided, the sum of
      positive weights for each unique y value, for each unique value in x.
      If y tensor is not provided, value is None.
    counts_per_x: if y is provided, counts of each of the unique values in x,
      otherwise, None.
  """

  if isinstance(y_input, tf.SparseTensor):
    # This is a multi-label/multi-task problem.

    # All cooccurrences tensors will have the same size to allow merging.
    # This should not influence the value of MI.
    max_y_value = tf.cast(
        tf.reduce_max(input_tensor=tf.sparse.to_dense(y_input)), tf.int64
    )
    # To handle this case, we densify the label with a padding value to compute
    # the cooccurences treating each "column" in the multivariant label as a
    # distinct Y label. Then we reduce by summing.
    result = tf.map_fn(
        lambda y_column: _compute_weighted_counts(  # pylint: disable=g-long-lambda
            x_input, y_column, weights_input, filter_regex, max_y_value
        ),
        # Set the default dense value to -1 to avoid confusing missing labels
        # with class 0.
        # Per segment_sum docs: If the given segment ID i is negative, the value
        # is dropped and will not be added to the sum of the segment.
        tf.transpose(
            tf.sparse.to_dense(
                tf.cast(y_input, tf.int64), default_value=_MISSING_MASK_VALUE
            )
        ),
        fn_output_signature={
            'unique_x_values': tf.TensorSpec(
                shape=[None],
                dtype=x_input.dtype,
            ),
            'summed_weights_per_x': tf.TensorSpec(
                shape=[None],
                dtype=tf.float32,
            ),
            'summed_positive_per_x_and_y': tf.TensorSpec(
                shape=[None, None],
                dtype=tf.float32,
            ),
            'unique_count': tf.TensorSpec(
                shape=[None],
                dtype=tf.int64,
            ),
        },
    )
    # summed_positive_per_x_and_y requires summing across the N "sub-labels".
    # For the other values, they should be a constant for each invocation and
    # thus we simply take the first: these values are computed on the X tensor
    # only, and won't depend on y_column (on which we're iterating).
    result['summed_positive_per_x_and_y'] = tf.reduce_sum(
        result['summed_positive_per_x_and_y'], axis=0
    )
    result['unique_x_values'] = result['unique_x_values'][0]
    result['summed_weights_per_x'] = result['summed_weights_per_x'][0]
    result['unique_count'] = result['unique_count'][0]
  else:
    max_y_value = tf.cast(tf.reduce_max(input_tensor=y_input), tf.int64)
    result = _compute_weighted_counts(
        x_input, y_input, weights_input, filter_regex, max_y_value
    )
  reduced_batch = ReducedBatchWeightedCounts(
      unique_x=result['unique_x_values'],
      summed_weights_per_x=result['summed_weights_per_x'],
      summed_positive_per_x_and_y=result['summed_positive_per_x_and_y'],
      counts_per_x=result['unique_count'],
  )
  # Add a sentinel token tracking the full distribution of y values.
  if extend_with_sentinel_counts:
    reduced_batch = extend_reduced_batch_with_y_counts(
        reduced_batch, y_input, weights_input
    )
  return reduced_batch
Performs batch-wise reduction to produce weighted co-occurrences.

  Computes the weighted co-occurrence of each feature value in x, for each value
  in the range [0, max(y)). If extend_with_sentinel_counts is true, the return
  value will include an additional sentinel token (not in the true vocabulary)
  that is used to accumulate the global distribution of y values.

  Args:
    x_input: Input `Tensor` or `CompositeTensor`.
    y_input: Integer `Tensor` or `SparseTensor` (in case of
      multi-label/multi-task) with which to compute the co-occurrence with
      x_input.
    weights_input: (Optional) Weights input `Tensor`.
    extend_with_sentinel_counts: If True, the reduced batch will be extended
      a sentinel value that accumlate the total distribution of y values. Should
      be True except when called recursively with the sentinel value as input.
    filter_regex: (Optional) Regex that matches tokens that have to be filtered
      out. Can only be specified if `x_input` has string dtype.

  Returns:
    a namedtuple of...
    unique_x_values: the unique values in x
    summed_weights_per_x: sum of the weights for each unique value in x
    summed_positive_per_x_and_y: If tensor y is provided, the sum of
      positive weights for each unique y value, for each unique value in x.
      If y tensor is not provided, value is None.
    counts_per_x: if y is provided, counts of each of the unique values in x,
      otherwise, None.
  
tensorflow_transform/tf_utils.py