code: 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))
  
file_path: jraph/jraph/examples/lstm.py
imports: import collections
  from absl import app
  import haiku
  import jax
  import jax.numpy
  import jax.tree_util
  import jraph
  import numpy
label_desc: Use image node_id 3 for calling a global function with example usage: main(_) without return types