code: def get_dummy_raw_datasets(
      dataset_length,
  ) -> Dict[str, tf.data.Dataset]:
      """Returns dummy datasets, mocking tfds.DatasetBuilder.as_dataset()."""
  
      # The dummy graph.
      num_nodes = 3
      num_edges = 4
      dummy_graph = {
          "edge_feat": tf.zeros((num_edges, 3), dtype=tf.float32),
          "edge_index": tf.zeros((num_edges, 2), dtype=tf.int64),
          "labels": tf.ones((128,), dtype=tf.float32),
          "node_feat": tf.zeros((num_nodes, 9), dtype=tf.float32),
          "num_edges": tf.expand_dims(num_edges, axis=0),
          "num_nodes": tf.expand_dims(num_nodes, axis=0),
      }
      dummy_graph_spec = {
          "edge_feat": tf.TensorSpec(shape=(None, 3), dtype=tf.float32),
          "edge_index": tf.TensorSpec(shape=(None, 2), dtype=tf.int64),
          "labels": tf.TensorSpec(shape=(128,), dtype=tf.float32),
          "node_feat": tf.TensorSpec(shape=(None, 9), dtype=tf.float32),
          "num_edges": tf.TensorSpec(shape=(None,), dtype=tf.int64),
          "num_nodes": tf.TensorSpec(shape=(None,), dtype=tf.int64),
      }
  
      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=dummy_graph_spec
          )
      return datasets
  
imports: import functools
  import pathlib
  import tempfile
  from typing import Dict, Optional
  import warnings
  from absl.testing import absltest
  from absl.testing import parameterized
  import flax
  from flax.training import train_state
  import jax
  from jax import numpy
  import jraph
  import tensorflow
  import tensorflow_datasets
  import numpy
  from configs import default
  from configs import test
  import input_pipeline
  import train
label_desc: Use image node_id 2 for calling a global function with example usage: get_dummy_raw_datasets(dataset_length) and returns: datasets