code: def _get_interaction_network(graphs_tuple):
      update_node_fn = lambda n, r: jnp.concatenate((n, r), axis=-1)
      update_edge_fn = lambda e, s, r: jnp.concatenate(
          (e, s, r), axis=-1
      )
      out = models.InteractionNetwork(update_edge_fn, update_node_fn)(
          graphs_tuple
      )
      nodes, edges, receivers, senders, _, _, _ = graphs_tuple
      expected_edges = jnp.concatenate(
          (edges, nodes[senders], nodes[receivers]), axis=-1
      )
      aggregated_nodes = utils.segment_sum(
          expected_edges,
          receivers,
          num_segments=len(graphs_tuple.nodes),
      )
      expected_nodes = jnp.concatenate(
          (nodes, aggregated_nodes), axis=-1
      )
      expected_out = graphs_tuple._replace(
          edges=expected_edges, nodes=expected_nodes
      )
      return out, expected_out
  
imports: from absl.testing import absltest
  from absl.testing import parameterized
  import jax
  import jax.numpy
  import jax.tree_util
  from jraph._src import graph
  from jraph._src import models
  from jraph._src import utils
  import numpy
label_desc: Use image node_id 9 for calling a global function with example usage: _get_interaction_network(graphs_tuple) and returns: out,  expected_out