code: def _get_multi_head_attention_graph_network(graphs_tuple):
      # Our test update functions are just identity functions.
      update_node_fn = lambda n, se, re, g: n
      update_global_fn = lambda gn, ge, g: g
  
      # With multi-head attention we have to return multiple edge features.
      # Here we define 3 heads, all with the same message.
      def update_edge_fn(e, unused_sn, unused_rn, unused_g):
          return tree.tree_map(lambda e_: jnp.stack([e_, e_, e_]), e)
  
      # Our attention logits are just the sum of the edge features of each head.
      def attention_logit_fn(e, unused_sn, unused_rn, unused_g):
          return tree.tree_map(lambda e_: jnp.sum(e_, axis=-1), e)
  
      # For multi-head attention we need a custom apply attention function.
      # In this we return the first edge.
      def attention_reduce_fn(e, unused_w):
          return tree.tree_map(lambda e_: e_[0], e)
  
      net = models.GraphNetwork(
          jax.vmap(update_edge_fn),
          jax.vmap(update_node_fn),
          update_global_fn,
          attention_logit_fn=jax.vmap(attention_logit_fn),
          attention_reduce_fn=jax.vmap(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 8 for calling a global function with example usage: _get_multi_head_attention_graph_network(graphs_tuple) and returns: net, tree, tree, tree