code: def get_dummy_datasets(
      dataset_length: int, batch_size: Optional[int] = None
  ) -> Dict[str, tf.data.Dataset]:
      """Returns dummy datasets, mocking input_pipeline.get_datasets()."""
  
      datasets = get_dummy_raw_datasets(dataset_length)
  
      # Construct the GraphsTuple converter function.
      convert_to_graphs_tuple_fn = functools.partial(
          input_pipeline.convert_to_graphs_tuple,
          add_virtual_node=True,
          add_undirected_edges=True,
          add_self_loops=True,
      )
  
      # Process each split separately.
      for split_name in datasets:
          # Convert to GraphsTuple.
          datasets[split_name] = datasets[split_name].map(
              convert_to_graphs_tuple_fn,
              num_parallel_calls=tf.data.AUTOTUNE,
              deterministic=True,
          )
  
      # If batch size is None, do not batch.
      if batch_size is not None:
          budget = (
              input_pipeline.estimate_padding_budget_for_batch_size(
                  datasets["train"], batch_size, num_estimation_graphs=1
              )
          )
  
          # Pad an example graph to see what the output shapes will be.
          # We will use this shape information when creating the tf.data.Dataset.
          example_graph = next(datasets["train"].as_numpy_iterator())
          example_padded_graph = jraph.pad_with_graphs(
              example_graph, *budget
          )
          padded_graphs_spec = input_pipeline.specs_from_graphs_tuple(
              example_padded_graph
          )
  
          # Batch and pad each split separately.
          for split, dataset_split in datasets.items():
              batching_fn = functools.partial(
                  jraph.dynamically_batch,
                  graphs_tuple_iterator=iter(dataset_split),
                  n_node=budget.n_node,
                  n_edge=budget.n_edge,
                  n_graph=budget.n_graph,
              )
              datasets[split] = tf.data.Dataset.from_generator(
                  batching_fn, output_signature=padded_graphs_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 3 for calling a global function with example usage: get_dummy_datasets(dataset_length, batch_size) and returns: datasets