code: def network_definition(graph: jraph.GraphsTuple) -> jraph.ArrayTree:
      """Implements the GCN from Kipf et al https://arxiv.org/pdf/1609.02907.pdf.
  
         A' = D^{-0.5} A D^{-0.5}
         Z = f(X, A') = A' relu(A' X W_0) W_1
  
      Args:
        graph: GraphsTuple the network processes.
  
      Returns:
        processed nodes.
      """
      gn = jraph.GraphConvolution(
          update_node_fn=hk.Linear(5, with_bias=False),
          add_self_edges=True,
      )
      graph = gn(graph)
      graph = graph._replace(nodes=jax.nn.relu(graph.nodes))
      gn = jraph.GraphConvolution(
          update_node_fn=hk.Linear(2, with_bias=False)
      )
      graph = gn(graph)
      return graph.nodes
  
file_path: jraph/jraph/examples/zacharys_karate_club.py
imports: import logging
  from absl import app
  import haiku
  import jax
  import jax.numpy
  import jraph
  import optax
label_desc: Use image node_id 3 for calling a global function with example usage: network_definition(graph) and returns: graph