def unsorted_segment_logsumexp(data: tf.Tensor, segment_ids: tf.Tensor, num_segments: Union[int, tf.Tensor]) -> tf.Tensor:
    """Computes the logarithmic sum of exponents along segments of a tensor
    like other operators from tf.math.unsorted_segment_* family.

    Args:
        data:           tf.Tensor,  shape = [...] + data_dims,
        segment_ids:    tf.Tensor,  shape = [...], dtype = tf.int32
        num_segments:   tf.Tensor,  shape = [], dtype = tf.int32

    Returns:            tf.Tensor,  shape = [num_segments] + data_dims, for the same type as data
    """
    data_max = tf.math.unsorted_segment_max(data=data, segment_ids=segment_ids, num_segments=num_segments)
    data_normed = data - tf.gather(params=data_max, indices=segment_ids)
    output = data_max + tf.math.log(
        tf.math.unsorted_segment_sum(
            data=tf.exp(data_normed),
            segment_ids=segment_ids,
            num_segments=num_segments,
        )
    )
    return output
Computes the logarithmic sum of exponents along segments of a tensor
    like other operators from tf.math.unsorted_segment_* family.

    Args:
        data:           tf.Tensor,  shape = [...] + data_dims,
        segment_ids:    tf.Tensor,  shape = [...], dtype = tf.int32
        num_segments:   tf.Tensor,  shape = [], dtype = tf.int32

    Returns:            tf.Tensor,  shape = [num_segments] + data_dims, for the same type as data
    
tensorflow_asr/losses/impl/ctc_tpu.py