code: def _unsorted_segment_sum(data, segment_ids, num_segments, name=None):
      data = _convert_to_tensor(data)
      del name
      if not JAX_MODE:
          raise NotImplementedError
      sums = np.zeros(num_segments)
      return sums.at[segment_ids].add(data)
  
imports: import collections
  import functools
  import numpy
  import numpy
  from tensorflow_probability.python.internal.backend.numpy import _utils
  from tensorflow_probability.python.internal.backend.numpy.numpy_array import _reverse
  from tensorflow_probability.python.internal.backend.numpy.numpy_array import one_hot
  from tensorflow_probability.python.internal.backend.numpy.ops import _convert_to_tensor
label_desc: Use image node_id 23 for calling a global function with example usage: _unsorted_segment_sum(data, segment_ids, num_segments, name) and returns: sums