def main(_):

  network = hk.without_apply_rng(hk.transform(network_definition))
  input_graph = get_random_graph()
  params = network.init(jax.random.PRNGKey(42), input_graph)
  output_graph = network.apply(params, input_graph)
  print(tree.tree_map(lambda x: x.shape, output_graph))