code: def _unsorted_segment_softmax(
      data, segment_ids, num_segments, name="unsorted_segment_softmax"
  ):
      """Performs an elementwise softmax operation along segments of a tensor.
  
      The input parameters are analogous to `tf.math.unsorted_segment_sum`. It
      produces an output of the same shape as the input data, after performing an
      elementwise sofmax operation between all of the rows with common segment id.
  
      Args:
        data: A tensor with at least one dimension.
        segment_ids: A tensor of indices segmenting `data` across the first
          dimension.
        num_segments: A scalar tensor indicating the number of segments. It should
          be at least `max(segment_ids) + 1`.
        name: A name for the operation (optional).
  
      Returns:
        A tensor with the same shape as `data` after applying the softmax operation.
  
      """
      with tf.name_scope(name):
          segment_maxes = tf.math.unsorted_segment_max(
              data, segment_ids, num_segments
          )
          maxes = tf.gather(segment_maxes, segment_ids)
          # Possibly refactor to `tf.stop_gradient(maxes)` for better performance.
          data -= maxes
          exp_data = tf.exp(data)
          segment_sum_exp_data = tf.math.unsorted_segment_sum(
              exp_data, segment_ids, num_segments
          )
          sum_exp_data = tf.gather(segment_sum_exp_data, segment_ids)
          return exp_data / sum_exp_data
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from graph_nets import _base
  from graph_nets import blocks
  import tensorflow
label_desc: Use image node_id 4 for calling a global function with example usage: _unsorted_segment_softmax(data, segment_ids, num_segments, name) and returns: unknown