code: def _segment_sum(
      data, segment_ids, name=None
  ):  # pylint: disable=unused-argument
      indices = _segment_ids_to_np_indices(segment_ids)
      if JAX_MODE:
          base = np.zeros_like(np.array(data)[indices])
          return base.at[segment_ids].add(data)
      else:
          return np.add.reduceat(data, indices)
  
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 16 for calling a global function with example usage: _segment_sum(data, segment_ids, name) and returns: base, np