class_docstr: Class for linear evaluation.
  
  Attributes:
    representation_fn: A function that is passed a TrainState and a Batch, and
      returns a tuple of (representation, label, and mask) for that batch. label
      and mask are ignored.
    rng: A PRNG that will be used to preprocess the data and to initialize the
      linear probe.
    config: A ConfigDict.
    linear_probe_cls: A Flax Module that is used as the linear probe. It must
      have a constructor argument called `num_classes`. Other constructor
      arguments are passed through via `linear_probe_init_kwargs`.
    linear_probe_init_kwargs: Keyword arguments passed to the constructor of
      `linear_probe_cls`.
class_docstr_tok: ["Class","for","linear","evaluation",".","Attributes",":","representation_fn",":","A","function","that","is","passed","a","TrainState","and","a","Batch",",","and","returns","a","tuple","of","(","representation",",","label",",","and","mask",")","for","that","batch",".","label","and","mask","are","ignored",".","rng",":","A","PRNG","that","will","be","used","to","preprocess","the","data","and","to","initialize","the","linear","probe",".","config",":","A","ConfigDict",".","linear_probe_cls",":","A","Flax","Module","that","is","used","as","the","linear","probe",".","It","must","have","a","constructor","argument","called","`","num_classes","`",".","Other","constructor","arguments","are","passed","through","via","`","linear_probe_init_kwargs","`",".","linear_probe_init_kwargs",":","Keyword","arguments","passed","to","the","constructor","of","`","linear_probe_cls","`","."]
docstr: Lazy-loads given dataset.
docstr_tok: ["Lazy-loads","given","dataset","."]
code: def _get_dataset(
      self,
      ds_name: str,
      config: ml_collections.ConfigDict,
      rng: jnp.ndarray,
  ) -> dataset_utils.Dataset:
      """Lazy-loads given dataset."""
      try:
          return self._datasets[ds_name]
      except KeyError:
          dataset = train_utils.get_dataset(config, rng)
          return self._datasets.setdefault(ds_name, dataset)
  
imports: import functools
  import logging
  from typing import Any, Callable, Dict, List, Mapping, Optional, Text, Tuple, Type, Union
  from absl import logging
  from clu import metric_writers
  import flax
  from flax import jax_utils
  import flax.linen
  import jax
  from jax.example_libraries.optimizers import clip_grads
  import jax.numpy
  import ml_collections
  import numpy
  import optax
  from scenic.dataset_lib import dataset_utils
  from scenic.model_lib.base_models import classification_model
  from scenic.model_lib.base_models import model_utils
  from scenic.train_lib import classification_trainer
  from scenic.train_lib import lr_schedules
  from scenic.train_lib import optimizers
  from scenic.train_lib import train_utils
label_desc: Use image node_id 2 for calling the LinearEvaluator obj's underlying member method code with example usage: obj._get_dataset(ds_name, config, rng) and returns: self, self