def get_random_graph() -> jraph.GraphsTuple:
  return jraph.GraphsTuple(
      n_node=np.asarray([NUM_NODES]),
      n_edge=np.asarray([NUM_EDGES]),
      nodes=np.random.normal(size=[NUM_NODES, EMBEDDING_SIZE]),
      edges=np.random.normal(size=[NUM_EDGES, EMBEDDING_SIZE]),
      globals=None,
      senders=np.random.randint(0, NUM_NODES, [NUM_EDGES]),
      receivers=np.random.randint(0, NUM_NODES, [NUM_EDGES]))