code: def accuracy(
      prediction: jnp.ndarray, target: jnp.ndarray
  ) -> jnp.float32:
      predicted_class = jnp.argmax(prediction, axis=1)
      return jnp.mean(predicted_class == target)
  
file_path: evojax/evojax/task/mnist.py
imports: import sys
  import numpy
  from typing import Tuple
  import jax
  import jax.numpy
  from jax import random
  from flax.struct import dataclass
  from evojax.task.base import VectorizedTask
  from evojax.task.base import TaskState
label_desc: Use image node_id 3 for calling a global function with example usage: accuracy(prediction, target) and returns: jnp