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 over multiple datasets.
  
  Args:
    repr_train_state: The train state that should be passed in as the first
      argument to `representation_fn`.
    datasets: A dictionary of datasets to evaluate on. The keys are names of
      the dataset that are used as the keys of the output dictionary. The
      values are configurations for the linear probe.
    writer: A metric writer.
    repr_step: The training step of the representation model being evaluated.
  
  Returns:
    A dictionary where the keys are the same as `datasets` and the values are
    lists of metrics computed over each test batch of that dataset.
docstr_tok: ["Computes","linear","evaluation","metrics","over","multiple","datasets",".","Args",":","repr_train_state",":","The","train","state","that","should","be","passed","in","as","the","first","argument","to","`","representation_fn","`",".","datasets",":","A","dictionary","of","datasets","to","evaluate","on",".","The","keys","are","names","of","the","dataset","that","are","used","as","the","keys","of","the","output","dictionary",".","The","values","are","configurations","for","the","linear","probe",".","writer",":","A","metric","writer",".","repr_step",":","The","training","step","of","the","representation","model","being","evaluated",".","Returns",":","A","dictionary","where","the","keys","are","the","same","as","`","datasets","`","and","the","values","are","lists","of","metrics","computed","over","each","test","batch","of","that","dataset","."]
code: def run_all(
      self,
      repr_train_state: Any,
      datasets: Dict[Text, ml_collections.ConfigDict],
      writer: metric_writers.MetricWriter,
      repr_step: int,
  ) -> Dict[Text, List[Metric]]:
      """Computes linear evaluation metrics over multiple datasets.
  
      Args:
        repr_train_state: The train state that should be passed in as the first
          argument to `representation_fn`.
        datasets: A dictionary of datasets to evaluate on. The keys are names of
          the dataset that are used as the keys of the output dictionary. The
          values are configurations for the linear probe.
        writer: A metric writer.
        repr_step: The training step of the representation model being evaluated.
  
      Returns:
        A dictionary where the keys are the same as `datasets` and the values are
        lists of metrics computed over each test batch of that dataset.
      """
  
      # Prepare the representation function given the current state of training.
      def representation_fn(batch: Batch) -> jnp.ndarray:
          return self.representation_fn(
              train_state=jax_utils.unreplicate(repr_train_state),
              batch=batch,
          )[
              0
          ]  # pytype: disable=wrong-keyword-args
  
      results = {}
      for ds_name, ds_cfg in datasets.items():
          logging.info("[linear_eval][%s]", ds_name)
          self.rng, ds_rng = jax.random.split(self.rng)
          results[ds_name] = self.run_one_dataset(
              ds_name,
              ds_cfg,
              representation_fn,
              ds_rng,
              writer,
              repr_step,
          )
  
      return results
  
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 6 for calling the LinearEvaluator obj's underlying member method code with example usage: obj.run_all(repr_train_state, datasets, writer, repr_step) and returns: results, self