code: def _load_dataset(name, split="train"):
      return tfds.load(
          name=name,
          split=split,
          data_dir=_TESTDATA,
          shuffle_files=False,
      )
  
imports: import collections
  import os
  from unittest import mock
  import gin
  import numpy
  from t5.data import assert_dataset
  from t5.data import preprocessors
  import tensorflow
  import tensorflow_datasets
  from trax.data import inputs
  from trax.data import tf_inputs
label_desc: Use image node_id 2 for calling a global function with example usage: _load_dataset(name, split) and returns: tfds