code: def _get_attention_graph_network(graphs_tuple):
      # Our test update functions are just identity functions.
      update_node_fn = lambda n, se, re, g: n
      update_edge_fn = lambda e, sn, rn, g: e
      update_global_fn = lambda gn, ge, g: g
      # Our attention logits are just one in this case.
      attention_logit_fn = lambda e, sn, rn, g: jnp.array(1.0)
      # We use a custom apply function here, which just returns the edge unchanged.
      attention_reduce_fn = lambda e, w: e
      net = models.GraphNetwork(
          update_edge_fn,
          update_node_fn,
          update_global_fn,
          attention_logit_fn=attention_logit_fn,
          attention_reduce_fn=attention_reduce_fn,
      )
      return net(graphs_tuple)
  
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 6 for calling a global function with example usage: _get_attention_graph_network(graphs_tuple) and returns: net