def _compute_summed_positive_per_x_and_y(
    weights, dummy_index, max_x_idx, max_y_value
):
  """Computes a segment sum to retrieve weighted co-occurrences of X and Y.

  Args:
    weights: The example weights, or a ones_like Tensor (=unweighted).
    dummy_index: An index the size of X containing shifted Y values. See
      cl/251659477.
    max_x_idx: The total number of distinct elements seen in the batch.
    max_y_value: The maximum id for Y values. Basically, the number of distinct
      values in Y.

  Returns:
    A 2D Tensor of shape [x_dim, y_dim] containing the weighted cooccurences of
    each Y value for each X value.
  """
  summed_positive_per_x_and_y = tf.cast(
      tf.math.unsorted_segment_sum(
          weights, dummy_index, max_x_idx * (max_y_value + 1)
      ),
      dtype=tf.float32,
  )
  summed_positive_per_x_and_y = tf.reshape(
      summed_positive_per_x_and_y, [max_x_idx, max_y_value + 1]
  )
  return summed_positive_per_x_and_y
Computes a segment sum to retrieve weighted co-occurrences of X and Y.

  Args:
    weights: The example weights, or a ones_like Tensor (=unweighted).
    dummy_index: An index the size of X containing shifted Y values. See
      cl/251659477.
    max_x_idx: The total number of distinct elements seen in the batch.
    max_y_value: The maximum id for Y values. Basically, the number of distinct
      values in Y.

  Returns:
    A 2D Tensor of shape [x_dim, y_dim] containing the weighted cooccurences of
    each Y value for each X value.
  
tensorflow_transform/tf_utils.py