code: def network_definition(
      graph: jraph.GraphsTuple, num_message_passing_steps: int = 1
  ) -> jraph.ArrayTree:
      """Defines a graph neural network.
  
      Args:
        graph: Graphstuple the network processes.
        num_message_passing_steps: number of message passing steps.
  
      Returns:
        globals.
      """
  
      @jax.vmap
      def update_fn(*args):
          size = args[0].shape[-1]
          return hk.nets.MLP([size, size])(
              jnp.concatenate(args, axis=-1)
          )
  
      for _ in range(num_message_passing_steps):
          gn = jraph.DeepSets(
              update_node_fn=update_fn,
              update_global_fn=update_fn,
              aggregate_nodes_for_globals_fn=jraph.segment_mean,
          )
          graph = gn(graph)
  
      return hk.Linear(graph.globals.shape[-1])(graph.globals)
  
file_path: jraph/jraph/examples/e_voting.py
imports: import collections
  import logging
  import random
  from absl import app
  import haiku
  import jax
  import jax.numpy
  import jraph
  import numpy
  import optax
label_desc: Use image node_id 2 for calling a global function with example usage: network_definition(graph, num_message_passing_steps) and returns: hk, hk