def _compute_weighted_counts(
    x_input, y_input, weights_input, filter_regex, max_y_value
):
  """Computes weighted counts of feature/label values to support AMI.

  Args:
    x_input: Input `Tensor` or `CompositeTensor`.
    y_input: Integer `Tensor` with which to compute the co-occurrence with
      x_input.
    weights_input: (Optional) Weights input `Tensor`.
    filter_regex: (Optional) Regex that matches tokens that have to be filtered
      out. Can only be specified if `x_input` has string dtype.
    max_y_value: The maximum index for labels (basically, the number of labels).

  Returns:
    A dictionary containing:
      - summed_positive_per_x_and_y: The occurrences of each X value, for each
        label value.
      - unique_x_values: All the X values in the vocabulary.
      - summed_weights_per_x: The sum of weights for each value in
        unique_x_values.
      - unique_count: The count of the values seen for X
        (basically, len(unique_x_values)).
  """
  # TODO(b/297854080): Evaluate the redundant calls.
  x, y, weights = _preprocess_tensors_for_cooccurences(
      x_input, y_input, weights_input, filter_regex
  )
  # TODO(b/297854080): Evaluate the redundant calls.
  vocabulary = _compute_vocabulary_values(x, weights)

  # Get a mask for the missing labels
  missing_labels_mask = tf.equal(y, _MISSING_MASK_VALUE)
  # Ultimately we want to compute the summed weights per x and y, as a
  # 2D [x_dim, y_dim] tensor. To accomplish this with a single call to
  # tf.math.unsorted_segment_sum, we flatten the [x_dim, y_dim] into a single
  # x*y dim vector and then re-shape back into 2d matrix form after computing
  # the segmented sum.
  flattened_index = (max_y_value + 1) * vocabulary.unique_idx + y
  # and apply it to the flattened index to keep missing values at -1.
  flattened_index = tf.where(
      missing_labels_mask,
      tf.constant(_MISSING_MASK_VALUE, tf.int64),
      flattened_index,
  )

  summed_positive_per_x_and_y = _compute_summed_positive_per_x_and_y(
      weights, flattened_index, vocabulary.max_x_idx, max_y_value
  )
  return {
      'summed_positive_per_x_and_y': tf.cast(
          summed_positive_per_x_and_y, tf.float32
      ),
      'unique_x_values': vocabulary.unique_x_values,
      'summed_weights_per_x': tf.cast(
          vocabulary.summed_weights_per_x, tf.float32
      ),
      'unique_count': tf.cast(vocabulary.unique_count, tf.int64),
  }
Computes weighted counts of feature/label values to support AMI.

  Args:
    x_input: Input `Tensor` or `CompositeTensor`.
    y_input: Integer `Tensor` with which to compute the co-occurrence with
      x_input.
    weights_input: (Optional) Weights input `Tensor`.
    filter_regex: (Optional) Regex that matches tokens that have to be filtered
      out. Can only be specified if `x_input` has string dtype.
    max_y_value: The maximum index for labels (basically, the number of labels).

  Returns:
    A dictionary containing:
      - summed_positive_per_x_and_y: The occurrences of each X value, for each
        label value.
      - unique_x_values: All the X values in the vocabulary.
      - summed_weights_per_x: The sum of weights for each value in
        unique_x_values.
      - unique_count: The count of the values seen for X
        (basically, len(unique_x_values)).
  
tensorflow_transform/tf_utils.py