python/dpu_utils/tf2utils/unsorted_segment_ops.py [30:39]:
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    max_per_segment = tf.math.unsorted_segment_max(
        data=logits, segment_ids=segment_ids, num_segments=num_segments
    )
    scattered_maxes = tf.gather(params=max_per_segment, indices=segment_ids)
    recentered_scores = logits - scattered_maxes
    exped_recentered_scores = tf.math.exp(recentered_scores)

    per_segment_sums = tf.math.unsorted_segment_sum(
        exped_recentered_scores, segment_ids, num_segments
    )
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -



python/dpu_utils/tf2utils/unsorted_segment_ops.py [51:60]:
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
    max_per_segment = tf.math.unsorted_segment_max(
        data=logits, segment_ids=segment_ids, num_segments=num_segments
    )
    scattered_maxes = tf.gather(params=max_per_segment, indices=segment_ids)
    recentered_scores = logits - scattered_maxes
    exped_recentered_scores = tf.math.exp(recentered_scores)

    per_segment_sums = tf.math.unsorted_segment_sum(
        exped_recentered_scores, segment_ids, num_segments
    )
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -