code: def InteractionNetwork(
      update_edge_fn: InteractionUpdateEdgeFn,
      update_node_fn: Union[
          InteractionUpdateNodeFn, InteractionUpdateNodeFnNoSentEdges
      ],
      aggregate_edges_for_nodes_fn: AggregateEdgesToNodesFn = utils.segment_sum,
      include_sent_messages_in_node_update: bool = False,
  ):
      """Returns a method that applies a configured InteractionNetwork.
  
      An interaction network computes interactions on the edges based on the
      previous edges features, and on the features of the nodes sending into those
      edges. It then updates the nodes based on the incoming updated edges.
      See https://arxiv.org/abs/1612.00222 for more details.
  
      This implementation adds an option not in https://arxiv.org/abs/1612.00222,
      which is to include edge features for which a node is a sender in the
      arguments to the node update function.
  
      Args:
        update_edge_fn: a function mapping a single edge update inputs to a single
          edge feature.
        update_node_fn: a function mapping a single node update input to a single
          node feature.
        aggregate_edges_for_nodes_fn: function used to aggregate messages to each
          node.
        include_sent_messages_in_node_update: pass edge features for which a node is
          a sender to the node update function.
      """
      # An InteractionNetwork is a GraphNetwork without globals features,
      # so we implement the InteractionNetwork as a configured GraphNetwork.
  
      # An InteractionNetwork edge function does not have global feature inputs,
      # so we filter the passed global argument in the GraphNetwork.
      wrapped_update_edge_fn = lambda e, s, r, g: update_edge_fn(
          e, s, r
      )
  
      # Similarly, we wrap the update_node_fn to ensure only the expected
      # arguments are passed to the Interaction net.
      if include_sent_messages_in_node_update:
          wrapped_update_node_fn = lambda n, s, r, g: update_node_fn(
              n, s, r
          )  # pytype: disable=wrong-arg-count
      else:
          wrapped_update_node_fn = lambda n, s, r, g: update_node_fn(
              n, r
          )  # pytype: disable=wrong-arg-count
      return GraphNetwork(
          update_edge_fn=wrapped_update_edge_fn,
          update_node_fn=wrapped_update_node_fn,
          aggregate_edges_for_nodes_fn=aggregate_edges_for_nodes_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 2 for calling a global function with example usage: InteractionNetwork(update_edge_fn, update_node_fn, aggregate_edges_for_nodes_fn, include_sent_messages_in_node_update) and returns: GraphNetwork