code: def _get_relation_network(graphs_tuple):
      edge_fn = lambda s, r: jnp.concatenate((s, r), axis=-1)
      global_fn = lambda e: e * 2
      out = models.RelationNetwork(edge_fn, global_fn)(graphs_tuple)
      expected_edges = jnp.concatenate(
          (
              graphs_tuple.nodes[graphs_tuple.senders],
              graphs_tuple.nodes[graphs_tuple.receivers],
          ),
          axis=-1,
      )
      num_graphs = len(graphs_tuple.n_edge)
      edge_gr_idx = jnp.repeat(
          jnp.arange(num_graphs),
          graphs_tuple.n_edge,
          total_repeat_length=graphs_tuple.edges.shape[0],
      )
      aggregated_edges = utils.segment_sum(
          expected_edges, edge_gr_idx, num_segments=num_graphs
      )
      expected_out = graphs_tuple._replace(
          edges=expected_edges, globals=aggregated_edges * 2
      )
      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 11 for calling a global function with example usage: _get_relation_network(graphs_tuple) and returns: out,  expected_out