code: def network_definition(
      graph: jraph.GraphsTuple, num_message_passing_steps: int = 5
  ) -> jraph.ArrayTree:
      """Defines a graph neural network.
  
      Args:
        graph: Graphstuple the network processes.
        num_message_passing_steps: number of message passing steps.
  
      Returns:
        Decoded nodes.
      """
      embedding = jraph.GraphMapFeatures(
          embed_edge_fn=jax.vmap(hk.Linear(output_size=16)),
          embed_node_fn=jax.vmap(hk.Linear(output_size=16)),
      )
      graph = embedding(graph)
  
      @jax.vmap
      @jraph.concatenated_args
      def update_fn(features):
          net = hk.Sequential(
              [
                  hk.Linear(10),
                  jax.nn.relu,
                  hk.Linear(10),
                  jax.nn.relu,
                  hk.Linear(10),
                  jax.nn.relu,
              ]
          )
          return net(features)
  
      for _ in range(num_message_passing_steps):
          gn = jraph.InteractionNetwork(
              update_edge_fn=update_fn,
              update_node_fn=update_fn,
              include_sent_messages_in_node_update=True,
          )
          graph = gn(graph)
  
      return hk.Linear(2)(graph.nodes)
  
file_path: jraph/jraph/examples/sat.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, net