code: def segment_sum(
      data: jnp.ndarray,
      segment_ids: jnp.ndarray,
      num_segments: Optional[int] = None,
      indices_are_sorted: bool = False,
      unique_indices: bool = False,
  ):
      """Computes the sum within segments of an array.
  
      Jraph alias to `jax.ops.segment_sum
      <https://jax.readthedocs.io/en/latest/_autosummary/jax.ops.segment_sum.html>`_.
      Note that other segment operations in jraph are not aliases, but are rather
      defined inside the package. Similar to TensorFlow's `segment_sum
      <https://www.tensorflow.org/api_docs/python/tf/math/segment_sum>`_.
  
      Args:
        data: an array with the values to be summed.
        segment_ids: an array with integer dtype that indicates the segments of
          `data` (along its leading axis) to be summed. Values can be repeated and
          need not be sorted. Values outside of the range [0, num_segments) are
          dropped and do not contribute to the sum.
        num_segments: optional, an int with nonnegative value indicating the number
          of segments. The default is set to be the minimum number of segments that
          would support all indices in ``segment_ids``, calculated as
          ``max(segment_ids) + 1``. Since `num_segments` determines the size of the
          output, a static value must be provided to use ``segment_sum`` in a
          ``jit``-compiled function.
        indices_are_sorted: whether ``segment_ids`` is known to be sorted.
        unique_indices: whether ``segment_ids`` is known to be free of duplicates.
  
      Returns:
        An array with shape :code:`(num_segments,) + data.shape[1:]` representing
        the segment sums.
  
      Examples:
        Simple 1D segment sum:
  
        >>> data = jnp.arange(5)
        >>> segment_ids = jnp.array([0, 0, 1, 1, 2])
        >>> segment_sum(data, segment_ids)
        DeviceArray([1, 5, 4], dtype=int32)
  
        Using JIT requires static `num_segments`:
  
        >>> from jax import jit
        >>> jit(segment_sum, static_argnums=2)(data, segment_ids, 3)
        DeviceArray([1, 5, 4], dtype=int32)
      """
      return jax.ops.segment_sum(
          data=data,
          segment_ids=segment_ids,
          num_segments=num_segments,
          indices_are_sorted=indices_are_sorted,
          unique_indices=unique_indices,
      )
  
imports: import functools
  from typing import Any, Callable, Generator, Iterable, Iterator, List, Mapping, Optional, Sequence, Union
  import jax
  from jax import lax
  import jax.numpy
  import jax.tree_util
  from jraph._src import graph
  import numpy
label_desc: Use image node_id 1 for calling a global function with example usage: segment_sum(data, segment_ids, num_segments, indices_are_sorted, unique_indices) and returns: jax