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: Evaluates a trained linear probe.
  
  Args:
    representation_fn: A function that given a batch, returns representations
      that are passed to the linear_probe_fn as input.
    linear_probe: The linear probe module instance.
    dataset: The dataset that has train_iter, valid_iter, meta_data, and
      optionally, test_iter.
    train_state: The train state of the trained LinearProbe.
    config: Configurations of evaluation, e.g., batch size.
    ds_name: The name of the training dataset.
    writer: A metric writer.
    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: ["Evaluates","a","trained","linear","probe",".","Args",":","representation_fn",":","A","function","that","given","a","batch",",","returns","representations","that","are","passed","to","the","linear_probe_fn","as","input",".","linear_probe",":","The","linear","probe","module","instance",".","dataset",":","The","dataset","that","has","train_iter",",","valid_iter",",","meta_data",",","and","optionally",",","test_iter",".","train_state",":","The","train","state","of","the","trained","LinearProbe",".","config",":","Configurations","of","evaluation",",","e.g.",",","batch","size",".","ds_name",":","The","name","of","the","training","dataset",".","writer",":","A","metric","writer",".","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 _eval(
      self,
      representation_fn: Callable[[Batch], jnp.ndarray],
      linear_probe: nn.Module,
      dataset: dataset_utils.Dataset,
      train_state: train_utils.TrainState,
      config: ml_collections.ConfigDict,
      ds_name: str,
      writer: metric_writers.MetricWriter,
      repr_step: int,
  ) -> List[Metric]:
      """Evaluates a trained linear probe.
  
      Args:
        representation_fn: A function that given a batch, returns representations
          that are passed to the linear_probe_fn as input.
        linear_probe: The linear probe module instance.
        dataset: The dataset that has train_iter, valid_iter, meta_data, and
          optionally, test_iter.
        train_state: The train state of the trained LinearProbe.
        config: Configurations of evaluation, e.g., batch size.
        ds_name: The name of the training dataset.
        writer: A metric writer.
        repr_step: The training step of the representation model being evaluated.
  
      Returns:
        A list of metrics computed over each batch of test data.
      """
      p_eval_step = jax.pmap(
          functools.partial(
              eval_step,
              representation_fn=representation_fn,
              linear_probe=linear_probe,
              metrics_fn=functools.partial(
                  classification_model.classification_metrics_function,
                  target_is_onehot=dataset.meta_data[
                      "target_is_onehot"
                  ],
              ),
          ),
          donate_argnums=(1,),
          axis_name="batch",
      )
      eval_metrics = []
      eval_batch_size = config.get("eval_batch_size", config.batch_size)
      total_eval_steps = int(
          np.ceil(
              dataset.meta_data["num_eval_examples"] / eval_batch_size
          )
      )
      for _ in range(total_eval_steps):
          batch = next(dataset.valid_iter)
          metrics = p_eval_step(train_state, batch)
          eval_metrics.append(train_utils.unreplicate_and_get(metrics))
  
      train_utils.log_eval_summary(
          step=repr_step,
          eval_metrics=eval_metrics,
          writer=writer,
          prefix=f"linear_eval/{ds_name}",
      )
  
      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 4 for calling the LinearEvaluator obj's underlying member method code with example usage: obj._eval(representation_fn, linear_probe, dataset, train_state, config, ds_name, writer, repr_step) and returns: eval_metrics