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: Computes linear evaluation metrics on one dataset.
  
  Args:
    ds_name: Name of the dataset, used for loading the data from cache dict.
    config: Configuration of the dataset, train and evaluation on.
    representation_fn: A function that given a batch, returns representations
      that are passed to the linear_probe_fn as input.
    rng: The JAX rng key.
    writer: A metric writer. Only needed if train summaries are to be written.
    repr_step: The training step of the representation model being evaluated.
  
  Returns:
    A list of metrics computed over each batch of test data.
docstr_tok: ["Computes","linear","evaluation","metrics","on","one","dataset",".","Args",":","ds_name",":","Name","of","the","dataset",",","used","for","loading","the","data","from","cache","dict",".","config",":","Configuration","of","the","dataset",",","train","and","evaluation","on",".","representation_fn",":","A","function","that","given","a","batch",",","returns","representations","that","are","passed","to","the","linear_probe_fn","as","input",".","rng",":","The","JAX","rng","key",".","writer",":","A","metric","writer",".","Only","needed","if","train","summaries","are","to","be","written",".","repr_step",":","The","training","step","of","the","representation","model","being","evaluated",".","Returns",":","A","list","of","metrics","computed","over","each","batch","of","test","data","."]
code: def run_one_dataset(
      self,
      ds_name: str,
      config: ml_collections.ConfigDict,
      representation_fn: Callable[[Batch], jnp.ndarray],
      rng: jnp.ndarray,
      writer: metric_writers.MetricWriter,
      repr_step: int,
  ) -> List[Metric]:
      """Computes linear evaluation metrics on one dataset.
  
      Args:
        ds_name: Name of the dataset, used for loading the data from cache dict.
        config: Configuration of the dataset, train and evaluation on.
        representation_fn: A function that given a batch, returns representations
          that are passed to the linear_probe_fn as input.
        rng: The JAX rng key.
        writer: A metric writer. Only needed if train summaries are to be written.
        repr_step: The training step of the representation model being evaluated.
  
      Returns:
        A list of metrics computed over each batch of test data.
      """
      data_rng, train_rng = jax.random.split(rng)
      dataset = self._get_dataset(ds_name, config, data_rng)
      linear_probe_head = self.linear_probe_cls(
          num_classes=dataset.meta_data["num_classes"],
          **self.linear_probe_init_kwargs,
      )
      logging.info(
          "[linear_eval]: Training linear probe for dataset %s", ds_name
      )
      train_state = self._train(
          representation_fn,
          linear_probe_head,
          dataset,
          train_rng,
          config,
          ds_name,
          writer,
          repr_step,
      )
      logging.info(
          "[linear_eval]: Evaluating linear probe for dataset %s",
          ds_name,
      )
      eval_metrics = self._eval(
          representation_fn,
          linear_probe_head,
          dataset,
          train_state,
          config,
          ds_name,
          writer,
          repr_step,
      )
      return eval_metrics
  
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 5 for calling the LinearEvaluator obj's underlying member method code with example usage: obj.run_one_dataset(ds_name, config, representation_fn, rng, writer, repr_step) and returns: eval_metrics