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: The main training loop.
  
  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.
    rng: JAX rng key.
    config: Configurations of the optimizer and learning rate scheduler.
    ds_name: The name of the training dataset.
    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:
    The train state of the trained linear probe.
docstr_tok: ["The","main","training","loop",".","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",".","rng",":","JAX","rng","key",".","config",":","Configurations","of","the","optimizer","and","learning","rate","scheduler",".","ds_name",":","The","name","of","the","training","dataset",".","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",":","The","train","state","of","the","trained","linear","probe","."]
code: def _train(
      self,
      representation_fn: Callable[[Batch], jnp.ndarray],
      linear_probe: nn.Module,
      dataset: dataset_utils.Dataset,
      rng: jnp.ndarray,
      config: ml_collections.ConfigDict,
      ds_name: str,
      writer: Optional[metric_writers.MetricWriter] = None,
      repr_step: int = 0,
  ) -> train_utils.TrainState:
      """The main training loop.
  
      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.
        rng: JAX rng key.
        config: Configurations of the optimizer and learning rate scheduler.
        ds_name: The name of the training dataset.
        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:
        The train state of the trained linear probe.
      """
      # Initialize model.
      input_shape = [1] + list(dataset.meta_data["input_shape"][1:])
      dummy_reprs = representation_fn(
          {
              "inputs": jnp.zeros(input_shape),
              "label": None,
              "batch_mask": None,
          }
      )
      model_state, params = flax.core.pop(
          linear_probe.init({"params": rng}, dummy_reprs), "params"
      )
  
      # Create optimizer.
      lr_fn = lr_schedules.get_learning_rate_fn(config)
      optimizer_config = optimizers.get_optax_optimizer_config(config)
      # If the config is already an optax-compatible config, better call directly:
      #   optimizers.get_optimizer(config.optimizer_configs, lr_fn)
      tx = optimizers.get_optimizer(optimizer_config, lr_fn)
      # We jit this, such that the arrays that are created are created on the same
      # device as the input is, in this case the CPU. Else they'd be on device[0].
      opt_state = jax.jit(tx.init, backend="cpu")(params)
  
      train_state = train_utils.TrainState(
          global_step=0,
          opt_state=opt_state,
          tx=tx,
          params=params,
          model_state=model_state,
      )
      train_state = jax_utils.replicate(train_state)
      p_train_step = jax.pmap(
          functools.partial(
              train_step,
              representation_fn=representation_fn,
              linear_probe=linear_probe,
              label_smoothing=config.get("label_smoothing"),
              max_grad_norm=config.get("max_grad_norm"),
          ),
          donate_argnums=(0, 1),
          axis_name="batch",
      )
      # Calculate the total number of training steps.
      total_steps, _ = train_utils.get_num_training_steps(
          config, dataset.meta_data
      )
      train_metrics = []
      extra_train_logs = []
      prefix = f"linear_eval_train/{ds_name}/step_{repr_step}"
      for step in range(total_steps):
          batch = next(dataset.train_iter)
          train_state, metrics = p_train_step(train_state, batch)
          train_metrics.append(metrics)
          # TODO(scenic-dev): Figure out how to get the lr from the optimizer.
          extra_train_logs.append(
              {f"{prefix}/learning_rate": lr_fn(step)}
          )
          if (step % 1000 == 0) or (step == total_steps - 1):
              logging.info("Linear probe trained for %d steps.", step)
          if (
              writer
              and self.config.get("log_train_summary_steps", 0) > 0
              and step % self.config.log_train_summary_steps == 0
          ):
              train_utils.log_train_summary(
                  step=step,
                  train_metrics=jax.tree_util.tree_map(
                      train_utils.unreplicate_and_get, train_metrics
                  ),
                  extra_training_logs=jax.tree_util.tree_map(
                      jax.device_get, extra_train_logs
                  ),
                  writer=writer,
                  prefix=prefix,
                  key_separator="/",
              )
              train_metrics = []
              extra_train_logs = []
      logging.info(
          "Linear probe training complete for dataset %s.", ds_name
      )
      return train_state
  
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 3 for calling the LinearEvaluator obj's underlying member method code with example usage: obj._train(representation_fn, linear_probe, dataset, rng, config, ds_name, writer, repr_step) and returns: train_state