code: def get_raw_datasets() -> Dict[str, tf.data.Dataset]:
      """Returns datasets as tf.data.Dataset, organized by split."""
      ds_builder = tfds.builder("ogbg_molpcba")
      ds_builder.download_and_prepare()
      ds_splits = ["train", "validation", "test"]
      datasets = {
          split: ds_builder.as_dataset(split=split)
          for split in ds_splits
      }
      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 1 for calling a global function with example usage: get_raw_datasets() and returns: datasets