code: def __init__(self, training: bool = True):
      dataset = load_dataset("mnist")
      dataset.set_format("np")
      X_train = np.stack(dataset["train"]["image"])
      y_train = dataset["train"]["label"]
      X_test = np.stack(dataset["test"]["image"])
      y_test = dataset["test"]["label"]
  
      if training:
          self.x = X_train
          self.y = y_train
      else:
          self.x = X_test
          self.y = y_test
  
file_path: elegy/examples/elegy/mnist_dataloader.py
imports: import os
  from dataclasses import dataclass
  from datetime import datetime
  from typing import Any, Generator, Mapping, Tuple
  import jax
  import jax.numpy
  import matplotlib.pyplot
  import numpy
  import optax
  import typer
  from datasets.load import load_dataset
  from tensorboardX.writer import SummaryWriter
  import elegy
label_desc: Use image node_id 1 to create a new MNIST object from inherited base classes: eg.data with example: obj = MNIST(training)