code: def main(_):
      config = get_config()
      train_state, logits_fn = create_train_state(
          _MODEL_PATH.value, config
      )
      train_state = jax_utils.replicate(train_state)
      p_logits_fn = jax.pmap(
          logits_fn, donate_argnums=(1,), axis_name="batch"
      )
      rng = jax.random.PRNGKey(config.rng_seed)
      data_rng, _ = jax.random.split(rng)
      dataset = train_utils.get_dataset(
          config, data_rng, dataset_service_address=None
      )
      total_eval_steps = int(
          np.ceil(
              dataset.meta_data["num_eval_examples"] / config.batch_size
          )
      )
      top1_correct_list = []
      for step in range(total_eval_steps):
          logging.info("Eval inference step: %d", step)
          batch = next(dataset.valid_iter)
          logits, labels, mask = p_logits_fn(train_state, batch)
          mask = np.array(jax_utils.unreplicate(mask)).astype(bool)
          logits = np.array(jax_utils.unreplicate(logits))[mask]
          labels = np.array(jax_utils.unreplicate(labels))[mask]
          top1_idx = jnp.argmax(logits, axis=-1)[..., None]
          top1_correct = jnp.take_along_axis(labels, top1_idx, axis=-1)
          top1_correct_list.append(top1_correct)
  
      top1_correct = np.concatenate(top1_correct_list, axis=0)
      acc = jnp.mean(top1_correct)
      logging.info("Classification acc: %f", acc)
  
file_path: scenic/scenic/projects/matvit/classification_eval_main.py
imports: import functools
  from typing import Callable, Dict, List, Optional, Tuple
  from absl import app
  from absl import flags
  from absl import logging
  from flax import jax_utils
  import flax.linen
  import jax
  import jax.numpy
  import jax.profiler
  import ml_collections
  import numpy
  from scenic.projects.matvit import matvit
  from scenic.train_lib import train_utils
label_desc: Use image node_id 4 for calling a global function with example usage: main(_) without return types