code: def net_fn(graph: jraph.GraphsTuple) -> jraph.GraphsTuple:
      """Graph net function."""
      # Add a global paramater for graph classification.
      graph = graph._replace(
          globals=jnp.zeros([graph.n_node.shape[0], 1])
      )
      embedder = jraph.GraphMapFeatures(
          hk.Linear(128), hk.Linear(128), hk.Linear(128)
      )
      net = jraph.GraphNetwork(
          update_node_fn=node_update_fn,
          update_edge_fn=edge_update_fn,
          update_global_fn=update_global_fn,
      )
      return net(embedder(graph))
  
imports: import functools
  import logging
  import pathlib
  import pickle
  from absl import app
  from absl import flags
  import haiku
  import jax
  import jax.numpy
  import jraph
  from jraph.ogb_examples import data_utils
  import optax
label_desc: Use image node_id 4 for calling a global function with example usage: net_fn(graph) and returns: net