code: def sharded_segment_sum(
      data, indices, num_segments, axis_index_groups
  ):
      """Segment sum over data on multiple devices."""
      device_segment_sum = utils.segment_sum(
          data, indices, num_segments
      )
      return jax.lax.psum(
          device_segment_sum,
          axis_name="i",
          axis_index_groups=axis_index_groups,
      )
  
file_path: jraph/jraph/experimental/sharded_graphnet.py
imports: import functools
  from typing import Callable, List, NamedTuple, Optional
  import jax
  import jax.numpy
  import jax.tree_util
  import jraph
  from jraph._src import graph
  from jraph._src import utils
  import numpy
label_desc: Use image node_id 3 for calling a global function with example usage: sharded_segment_sum(data, indices, num_segments, axis_index_groups) and returns: jax