code: def get_graphs_tuple_size(graph: jraph.GraphsTuple):
      """Returns the number of nodes, edges and graphs in a GraphsTuple."""
      return GraphsTupleSize(
          n_node=np.sum(graph.n_node),
          n_edge=np.sum(graph.n_edge),
          n_graph=np.shape(graph.n_node)[0],
      )
  
file_path: flax/examples/ogbg_molpcba/input_pipeline.py
imports: import functools
  from typing import Dict, NamedTuple
  import jraph
  import numpy
  import tensorflow
  import tensorflow_datasets
label_desc: Use image node_id 6 for calling a global function with example usage: get_graphs_tuple_size(graph) and returns: GraphsTupleSize