code: def unsorted_segment_softmax(x, indices, n_nodes=None):
      """
      Applies softmax along the segments of a Tensor. This operator is similar
      to the tf.math.segment_* operators, which apply a certain reduction to the
      segments. In this case, the output tensor is not reduced and maintains the
      same shape as the input.
      :param x: a Tensor. The softmax is applied along the first dimension.
      :param indices: a Tensor, indices to the segments.
      :param n_nodes: the number of unique segments in the indices. If `None`,
      n_nodes is calculated as the maximum entry in the indices plus 1.
      :return: a Tensor with the same shape as the input.
      """
      n_nodes = (
          tf.reduce_max(indices) + 1 if n_nodes is None else n_nodes
      )
      e_x = tf.exp(
          x
          - tf.gather(
              tf.math.unsorted_segment_max(x, indices, n_nodes), indices
          )
      )
      e_x /= tf.gather(
          tf.math.unsorted_segment_sum(e_x, indices, n_nodes) + 1e-9,
          indices,
      )
      return e_x
  
file_path: spektral/spektral/layers/ops/scatter.py
imports: import tensorflow
label_desc: Use image node_id 7 for calling a global function with example usage: unsorted_segment_softmax(x, indices, n_nodes) and returns: e_x