code: def get_dummy_datasets(dataset_length: int):
      """Returns a set of datasets of unbatched GraphsTuples."""
      # The dummy graph.
      num_nodes = 3
      num_edges = 4
      dummy_graph = jraph.GraphsTuple(
          n_node=tf.expand_dims(num_nodes, 0),
          n_edge=tf.expand_dims(num_edges, 0),
          senders=tf.zeros(num_edges, dtype=tf.int32),
          receivers=tf.ones(num_edges, dtype=tf.int32),
          nodes=tf.zeros((num_nodes, 9)),
          edges=tf.ones((num_edges, 3)),
          globals=tf.ones((1, 128), dtype=tf.int64),
      )
      graphs_spec = input_pipeline.specs_from_graphs_tuple(dummy_graph)
  
      # Yields a set of graphs for the current split.
      def get_dummy_graphs():
          for _ in range(dataset_length):
              yield dummy_graph
  
      datasets = {}
      for split in ["train", "validation", "test"]:
          datasets[split] = tf.data.Dataset.from_generator(
              get_dummy_graphs, output_signature=graphs_spec
          )
      return datasets
  
file_path: flax/examples/ogbg_molpcba/input_pipeline_test.py
imports: from absl.testing import absltest
  from absl.testing import parameterized
  import input_pipeline
  import jraph
  import tensorflow
label_desc: Use image node_id 1 for calling a global function with example usage: get_dummy_datasets(dataset_length) and returns: datasets