class_docstr: Class for few-shot evaluation.
class_docstr_tok: ["Class","for","few-shot","evaluation","."]
docstr: Lazy-loads given dataset.
docstr_tok: ["Lazy-loads","given","dataset","."]
code: def _get_dataset(self, dataset, train_split, test_split):
      """Lazy-loads given dataset."""
      key = (dataset, train_split, test_split)
      try:
          return self._datasets[key]
      except KeyError:
          train_ds = dataset_utils.get_data(
              dataset=dataset,
              split=train_split,
              batch_size=self.local_batch_size,
              preprocess_fn=pp_builder.get_preprocess_fn(self.pp_tr),
              repeats=1,
              cache="loaded",
              drop_remainder=False,
          )
          test_ds = dataset_utils.get_data(
              dataset=dataset,
              split=test_split,
              batch_size=self.local_batch_size,
              preprocess_fn=pp_builder.get_preprocess_fn(self.pp_te),
              repeats=1,
              cache="loaded",
              drop_remainder=False,
          )
          num_classes = (
              tfds.builder(dataset).info.features["label"].num_classes
          )
          return self._datasets.setdefault(
              key, (train_ds, test_ds, num_classes)
          )
  
imports: import functools
  from typing import Any, Optional, Tuple
  from absl import logging
  from clu import metric_writers
  from flax import jax_utils
  import jax
  import jax.numpy
  import numpy
  from scenic.dataset_lib import dataset_utils
  from scenic.dataset_lib.big_transfer import builder
  from scenic.train_lib import train_utils
  import tensorflow_datasets
label_desc: Use image node_id 2 for calling the FewShotEvaluator obj's underlying member method code with example usage: obj._get_dataset(dataset, train_split, test_split) and returns: self, self