code: def edge_update_fn(feats: jnp.ndarray) -> jnp.ndarray:
      """Edge update function for graph net."""
      net = hk.Sequential([hk.Linear(128), jax.nn.relu, hk.Linear(128)])
      return net(feats)
  
imports: import functools
  import logging
  import pathlib
  import pickle
  from typing import Iterator
  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 1 for calling a global function with example usage: edge_update_fn(feats) and returns: net