code: def RelationNetwork(
      update_edge_fn: Callable[
          [SenderFeatures, ReceiverFeatures], EdgeFeatures
      ],
      update_global_fn: Callable[[EdgeFeatures], NodeFeatures],
      aggregate_edges_for_globals_fn: AggregateEdgesToGlobalsFn = utils.segment_sum,
  ):
      """Returns a method that applies a Relation Network.
  
      See https://arxiv.org/abs/1706.01427 for more details.
  
      This implementation has one more argument, `aggregate_edges_for_globals_fn`,
      which changes how edge features are aggregated. The paper uses the default -
      `utils.segment_sum`.
  
      Args:
        update_edge_fn: function used to update the edges.
        update_global_fn: function used to update the globals.
        aggregate_edges_for_globals_fn: function used to aggregate the edges for the
          globals.
      """
      return GraphNetwork(
          update_edge_fn=lambda e, s, r, g: update_edge_fn(s, r),
          update_node_fn=None,
          update_global_fn=lambda n, e, g: update_global_fn(e),
          attention_logit_fn=None,
          aggregate_edges_for_globals_fn=aggregate_edges_for_globals_fn,
      )
  
file_path: jraph/jraph/_src/models.py
imports: import functools
  from typing import Any, Callable, Iterable, Mapping, Optional, Union
  import jax
  import jax.numpy
  import jax.tree_util
  from jraph._src import graph
  from jraph._src import utils
label_desc: Use image node_id 4 for calling a global function with example usage: RelationNetwork(update_edge_fn, update_global_fn, aggregate_edges_for_globals_fn) and returns: GraphNetwork