code: def segment_sum(data, segment_ids, num_segments=None):
      """Computes the sum along segment_ids along axis 0. If multiple segment_ids reference the same
      location their contributions add up.
      result[index, j, k, ...] = Σi... data[i, j, k,..] where index = segment_ids[i]
      This op is much better understood with visualization articulated in the following links and
      examples at the end of this docstring.
  
      https://www.tensorflow.org/api_docs/python/tf/math/unsorted_segment_sum
      https://caffe2.ai/docs/sparse-operations.html#null__unsorted-segment-reduction-ops
  
      Parameters
      ----------
      data : relay.Expr
          Input tensor. It can be of any type and multi-dimensional.
  
      segment_ids : relay.Expr
          A 1-D int32/int64 tensor containing the segment_ids of the rows to calculate the output
          sum upon. It defines a mapping from the zeroth dimension of data onto segment_ids. The
          segment_ids tensor should be the size of the first dimension, d0, with consecutive IDs
          in the range 0 to k, where k<d0. In particular, a segmentation of a matrix tensor is a
          mapping of rows to segments. This tensor doesn't need to be sorted.
  
      num_segments : int, optional
          An integer describing the shape of the zeroth dimension. If unspecified, it is calculated
          equivalent to the number of unique segment_ids.
  
      Returns
      -------
      result : relay.Expr
          Output tensor.
  
      Examples
      --------
      .. code-block:: python
  
          data = [[1, 2, 3, 4],
                  [4, -3, 2, -1],
                  [5, 6, 7, 8]]
  
          segment_ids = [0, 0, 1]
  
          relay.segment_sum(data, segment_ids) = [[5, -1, 5, 3],
                                                  [5, 6, 7, 8]]
  
          data = [[1, 2, 3, 4],
                  [4, -3, 2, -1],
                  [5, 6, 7, 8]]
  
          segment_ids = [2, 0, 0]
  
          num_segments = 3
  
          segment_sum(data, segment_ids, num_segments) = [[9, 3, 9, 7],
                                                          [0, 0, 0, 0],
                                                          [1, 2, 3, 4]]
      """
  
      one_tensor = cast_like(const([1]), segment_ids)
      if num_segments:
          if isinstance(num_segments, int):
              max_segments = const([num_segments])
              max_segments = cast_like(max_segments, segment_ids)
          else:
              max_segments = cast_like(num_segments, segment_ids)
      else:
          max_segments = _make.add(
              reshape(_make.max(segment_ids, [0], False, False), -1),
              one_tensor,
          )
  
      data_offrow_shape = strided_slice(
          _make.shape_of(data, "int32"), [1], [-1], slice_mode="size"
      )
      data_offrow_shape = cast_like(data_offrow_shape, max_segments)
      new_shape = _make.concatenate(
          Tuple([max_segments, data_offrow_shape]), 0
      )
      segment_ids_tiled_shape = _make.concatenate(
          Tuple([reverse(data_offrow_shape, 0), one_tensor]), 0
      )
      expanded_segment_ids = tile(segment_ids, segment_ids_tiled_shape)
      scatter_add_segment_ids = transpose(expanded_segment_ids)
      src = cast_like(_dyn_make.zeros(new_shape, "float64"), data)
      return scatter_elements(
          src, scatter_add_segment_ids, data, axis=0, reduction="add"
      )
  
imports: from ...tir import expr
  from ..expr import Constant, Expr, Tuple, TupleWrapper, const
  from .None import _make
  from .dyn import _make
  from .tensor import shape_of
label_desc: Use image node_id 43 for calling a global function with example usage: segment_sum(data, segment_ids, num_segments) and returns: scatter_elements