code: def main(_):
      def net_fn(graph: jraph.GraphsTuple):
          unf = jraph.concatenated_args(conway_mlp)
          net = jraph.InteractionNetwork(
              update_edge_fn=lambda e, n_s, n_r: n_s,
              update_node_fn=jax.vmap(unf),
          )
          return net(graph)
  
      net = hk.without_apply_rng(hk.transform(net_fn))
  
      cg = conway_graph(size=20)
      params = net.init(jax.random.PRNGKey(42), cg)
      for _ in range(100):
          time.sleep(0.05)
          cg = jax.jit(net.apply)(params, cg)
          display_graph(cg)
  
file_path: jraph/jraph/examples/game_of_life.py
imports: import time
  from absl import app
  import haiku
  import jax
  import jax.numpy
  import jraph
  import numpy
label_desc: Use image node_id 4 for calling a global function with example usage: main(_) and returns: net