code: def loss(prediction: jnp.ndarray, target: jnp.ndarray) -> jnp.float32:
      target = jax.nn.one_hot(target, 10)
      return -jnp.mean(jnp.sum(prediction * target, axis=1))
  
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 2 for calling a global function with example usage: loss(prediction, target) and returns: unknown