code: def test_accuracy(
      model: snt.Module,
      dataset: tf.data.Dataset,
  ) -> Dict[str, tf.Tensor]:
      """Computes accuracy on the test set."""
      correct, total = 0, 0
      for images, labels in dataset:
          preds = tf.argmax(model(images), axis=1)
          correct += tf.math.count_nonzero(
              tf.equal(preds, labels), dtype=tf.int32
          )
          total += tf.shape(labels)[0]
      accuracy = (correct / tf.cast(total, tf.int32)) * 100.0
      return {"accuracy": accuracy, "incorrect": total - correct}
  
file_path: sonnet/examples/simple_mnist.py
imports: from typing import Dict
  from absl import app
  import sonnet
  import tensorflow
  import tensorflow_datasets
label_desc: Use image node_id 4 for calling a global function with example usage: test_accuracy(model, dataset) and returns: dict