f73cdc286a3e125cf7019336621cb10370ebfd52,tensorflow_transform/tf_utils.py,,_reduce_vocabulary_inputs,#Any#Any#Any#,71

Before Change

    greater_assert = tf.Assert(
        tf.greater_equal(tf.reduce_min(input_tensor=labels), 0), [labels])
    with tf.control_dependencies([less_assert, greater_assert]):
      labels = tf.identity(labels)
    positive_weights = (
        tf.cast(labels, tf.float32) * tf.cast(weights, tf.float32))
    summed_positive_weights = tf.math.unsorted_segment_sum(
        positive_weights, unique.idx, tf.size(input=unique.y))

After Change

    max_label_value = tf.cast(tf.reduce_max(input_tensor=labels), tf.int32)
    one_hot_labels = tf.one_hot(labels, max_label_value + 1)
    broadcast_weights = tf.cast(
        tf.broadcast_to(tf.reshape(weights, (-1, 1)), tf.shape(one_hot_labels)),
        dtype=tf.float32)
    positive_weights = (tf.cast(one_hot_labels, tf.float32) * broadcast_weights)
    summed_positive_per_value_and_label = tf.math.unsorted_segment_sum(

Instances