code: def eval_batch(
      params: hk.Params,
      state: hk.State,
      batch: dataset.Batch,
  ) -> jax.Array:
      """Evaluates a batch."""
      logits, _ = forward.apply(
          params, state, None, batch, is_training=False
      )
      predicted_label = jnp.argmax(logits, axis=-1)
      correct = jnp.sum(jnp.equal(predicted_label, batch["labels"]))
      return correct.astype(jnp.float32)
  
imports: from collections.abc import Iterable, Mapping
  import contextlib
  import functools
  import timeit
  from typing import NamedTuple
  from absl import app
  from absl import flags
  from absl import logging
  import haiku
  from examples.imagenet import dataset
  import jax
  import jax.numpy
  import jmp
  import numpy
  import optax
  import tree
label_desc: Use image node_id 9 for calling a global function with example usage: eval_batch(params, state, batch) and returns: correct