class_docstr: Class for few-shot evaluation.
class_docstr_tok: ["Class","for","few-shot","evaluation","."]
docstr: Compute few-shot metrics on one dataset.
docstr_tok: ["Compute","few-shot","metrics","on","one","dataset","."]
code: def compute_fewshot_metrics(
      self, train_state, dataset, train_split, test_split
  ):
      """Compute few-shot metrics on one dataset."""
  
      (
          train_ds,
          test_ds,
          num_train_samples,
          num_test_samples,
          num_classes,
          is_one_hot,
      ) = self._get_dataset(
          dataset,
          train_split,
          test_split,
          self.config.get("num_train_examples"),
          self.config.get("num_test_examples"),
      )
      logging.info(
          "[fewshot][%s]: Precomputing train (%s)", dataset, train_split
      )
      repr_train, labels_train = self._get_repr(
          train_state, train_ds, num_train_samples
      )
      labels_train = jnp.squeeze(labels_train)
      logging.info(
          "[fewshot][%s]: Precomputing test (%s)", dataset, test_split
      )
      repr_test, labels_test = self._get_repr(
          train_state, test_ds, num_test_samples
      )
      labels_test = jnp.squeeze(labels_test)
      logging.info("[fewshot][%s]: Solving linear system", dataset)
      # Collect where we have samples of which classes.
      if is_one_hot:
          class_indices = [
              np.where(labels_train[:, cls_i] == 1)[0]
              for cls_i in range(num_classes)
          ]
      else:
          class_indices = [
              np.where(labels_train == cls_i)[0]
              for cls_i in range(num_classes)
          ]
  
      results = {}
      for shots in self.shots:
          all_idx = [indices[:shots] for indices in class_indices]
          all_idx = np.concatenate(all_idx, axis=0)
          x = repr_train[all_idx]
          y = labels_train[all_idx]
          for l2_reg in self.l2_regs:
              acc = _fewshot_acc_fn(
                  x,
                  y,
                  repr_test,
                  labels_test,
                  l2_reg,
                  num_classes,
                  is_one_hot,
              )
              results[shots, l2_reg] = np.array(acc)
      return results
  
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 4 for calling the FewShotEvaluatorVideo obj's underlying member method code with example usage: obj.compute_fewshot_metrics(train_state, dataset, train_split, test_split) and returns: results