class_docstr: Class for BERT few-shot evaluation.
class_docstr_tok: ["Class","for","BERT","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, num_samples):
      """Compute representation for the whole dataset."""
      pre_logits_list = []
      labels_list = []
      total_steps = int(np.ceil(num_samples / self.local_batch_size))
      for _ in range(1, total_steps + 1):
          batch = next(data)
          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]`.
          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]`."
              )
          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],
              )
          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]
          )
      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, Dict, List, Mapping, Optional, Tuple, Union
  from absl import logging
  from clu import metric_writers
  import flax
  from flax import jax_utils
  import flax.linen
  import jax
  import jax.numpy
  import jax.profiler
  import ml_collections
  import numpy
  from scenic.common_lib import debug_utils
  from scenic.train_lib import optimizers
  from scenic.train_lib import train_utils
  import scipy
  import sklearn.metrics
label_desc: Use image node_id 3 for calling the BERTFewShotEvaluator obj's underlying member method code with example usage: obj._get_repr(train_state, data, num_samples) and returns: pre_logits,  labels