class_docstr: Class for few-shot evaluation.
class_docstr_tok: ["Class","for","few-shot","evaluation","."]
docstr: Compute representation for the whole dataset.
docstr_tok: ["Compute","representation","for","the","whole","dataset","."]
code: def _get_repr(self, train_state, data):
      """Compute representation for the whole dataset."""
      pre_logits_list = []
      labels_list = []
      for batch in start_input_pipeline(
          data,
          pad=self.local_batch_size,
          backend=self.backend,
          devices=self.devices,
      ):
          pre_logits, labels, mask = self.repr_fn(train_state, batch)
          # We need to unreplicate the output of `lax.all_gather`.
          # Shapes at this point are:
          #   pre_logits: `[hosts, devices, global_batch, features]`.
          #   labels: `[hosts, devices, global_batch]`.
          #   mask: `[hosts, devices, global_batch]`.
          if self.backend == "pmap":
              pre_logits = jax_utils.unreplicate(pre_logits)
  
              if pre_logits.ndim != 3:
                  raise ValueError(
                      "Shape of the representations sent to the linear "
                      "fewshot should be `[num_devices, bs, features]`."
                  )
  
              mask = np.array(jax_utils.unreplicate(mask)).astype(bool)
              pre_logits_list.append(np.array(pre_logits)[mask])
              labels_list.append(
                  np.array(jax_utils.unreplicate(labels))[mask]
              )
          else:
              pre_logits = jax.device_get(pre_logits)
  
              if pre_logits.ndim != 2:
                  raise ValueError(
                      "Shape of the representations sent to the linear "
                      "fewshot should be `[bs, features]`."
                  )
  
              mask = np.array(jax.device_get(mask)).astype(bool)
              pre_logits_list.append(np.array(pre_logits)[mask])
              labels_list.append(np.array(jax.device_get(labels))[mask])
  
          if pre_logits.shape[-1] > 2048:
              logging.warning(
                  "The feature size for the representations is too large"
                  "(feature size = %d). This might cause severe slowdown "
                  "of solving the linear equation.",
                  pre_logits.shape[-1],
              )
  
      pre_logits = np.concatenate(pre_logits_list, axis=0)
      labels = np.concatenate(labels_list, axis=0)
      return pre_logits, labels
  
imports: import functools
  from typing import Any, Optional, Tuple
  from absl import logging
  from clu import metric_writers
  from flax import jax_utils
  import jax
  import jax.numpy
  import numpy
  from scenic.dataset_lib import dataset_utils
  from scenic.dataset_lib.big_transfer import builder
  from scenic.train_lib import train_utils
  import tensorflow_datasets
label_desc: Use image node_id 3 for calling the FewShotEvaluator obj's underlying member method code with example usage: obj._get_repr(train_state, data) and returns: pre_logits,  labels