class_docstr: MNIST classification task.
class_docstr_tok: ["MNIST","classification","task","."]
code: def reset(self, key: jnp.ndarray) -> State:
      return self._reset_fn(key)
  
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 the MNIST obj's underlying member method code with example usage: obj.reset(key) and returns: self