code: def _get_gat(graphs_tuple):
      # With multi-head attention we have to return multiple edge features.
      # Here we define 3 heads, all with the same message.
      def attention_query_fn(n):
          return tree.tree_map(
              lambda n_: jnp.stack([n_, n_, n_], axis=2), n
          )
  
      # Our attention logits 1 if a self edge
      def attention_logit_fn(s, r, e_):
          del e_
          return (s == r) * 1 + (s != r) * -1e10
  
      def node_update_fn(nodes):
          return jnp.mean(nodes, axis=2)
  
      net = models.GAT(
          attention_query_fn, attention_logit_fn, node_update_fn
      )
  
      # Cast nodes to floats since GAT will output floats from the softmax
      # attention.
      graphs_tuple = graphs_tuple._replace(
          nodes=jnp.array(graphs_tuple.nodes, jnp.float32)
      )
      return net(graphs_tuple), 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 13 for calling a global function with example usage: _get_gat(graphs_tuple) and returns: net,  graphs_tuple, tree, unknown, jnp