code: def conway_graph(size) -> jraph.GraphsTuple:
      """Returns a graph representing the game field of conway's game of life."""
      # Creates nodes: each node represents a cell in the game.
      n_node = size**2
      nodes = np.zeros((n_node, 1))
      node_indices = jnp.arange(n_node)
      # Creates edges, senders and receivers:
      # the senders represent the connections to the 8 neighboring fields.
      n_edge = 8 * n_node
      edges = jnp.zeros((n_edge, 1))
      senders = jnp.vstack(
          [
              node_indices - size - 1,
              node_indices - size,
              node_indices - size + 1,
              node_indices - 1,
              node_indices + 1,
              node_indices + size - 1,
              node_indices + size,
              node_indices + size + 1,
          ]
      )
      senders = senders.T.reshape(-1)
      senders = (senders + size**2) % size**2
      receivers = jnp.repeat(node_indices, 8)
      # Adds a glider to the game
      nodes[0, 0] = 1.0
      nodes[1, 0] = 1.0
      nodes[2, 0] = 1.0
      nodes[2 + size, 0] = 1.0
      nodes[1 + 2 * size, 0] = 1.0
      return jraph.GraphsTuple(
          n_node=jnp.array([n_node]),
          n_edge=jnp.array([n_edge]),
          nodes=jnp.asarray(nodes),
          edges=edges,
          globals=None,
          senders=senders,
          receivers=receivers,
      )
  
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 2 for calling a global function with example usage: conway_graph(size) and returns: jraph