code: def get_datasets(
      batch_size: int,
      add_virtual_node: bool = True,
      add_undirected_edges: bool = True,
      add_self_loops: bool = True,
  ) -> Dict[str, tf.data.Dataset]:
      """Returns datasets of batched GraphsTuples, organized by split."""
      if batch_size <= 1:
          raise ValueError(
              "Batch size must be > 1 to account for padding graphs."
          )
  
      # Obtain the original datasets.
      datasets = get_raw_datasets()
  
      # Construct the GraphsTuple converter function.
      convert_to_graphs_tuple_fn = functools.partial(
          convert_to_graphs_tuple,
          add_virtual_node=add_virtual_node,
          add_undirected_edges=add_undirected_edges,
          add_self_loops=add_self_loops,
      )
  
      # 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,
          )
  
      # Compute the padding budget for the requested batch size.
      budget = estimate_padding_budget_for_batch_size(
          datasets["train"], batch_size, num_estimation_graphs=100
      )
  
      # 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 = specs_from_graphs_tuple(example_padded_graph)
  
      # Process each split separately.
      for split_name, dataset_split in datasets.items():
          # Repeat and shuffle the training split.
          if split_name == "train":
              dataset_split = dataset_split.shuffle(
                  100, reshuffle_each_iteration=True
              )
              dataset_split = dataset_split.repeat()
  
          # Batch and pad each split.
          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,
          )
          dataset_split = tf.data.Dataset.from_generator(
              batching_fn, output_signature=padded_graphs_spec
          )
  
          # We cache the validation and test sets, since these are small.
          if split_name in ["validation", "test"]:
              dataset_split = dataset_split.cache()
  
          datasets[split_name] = dataset_split
      return datasets
  
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 2 for calling a global function with example usage: get_datasets(batch_size, add_virtual_node, add_undirected_edges, add__loops) and returns: datasets