code: def specs_from_graphs_tuple(graph: jraph.GraphsTuple):
      """Returns a tf.TensorSpec corresponding to this graph."""
  
      def get_tensor_spec(array: np.ndarray):
          shape = list(array.shape)
          dtype = array.dtype
          return tf.TensorSpec(shape=shape, dtype=dtype)
  
      specs = {}
      for field in [
          "nodes",
          "edges",
          "senders",
          "receivers",
          "globals",
          "n_node",
          "n_edge",
      ]:
          field_sample = getattr(graph, field)
          specs[field] = get_tensor_spec(field_sample)
      return jraph.GraphsTuple(**specs)
  
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 5 for calling a global function with example usage: specs_from_graphs_tuple(graph) and returns: jraph, tf