code: def main(
      debug: bool = False,
      eager: bool = False,
      logdir: str = "runs",
      steps_per_epoch: int = 200,
      epochs: int = 100,
      batch_size: int = 64,
  ):
      if debug:
          import debugpy
  
          print("Waiting for debugger...")
          debugpy.listen(5678)
          debugpy.wait_for_client()
  
      current_time = datetime.now().strftime("%b%d_%H-%M-%S")
      logdir = os.path.join(logdir, current_time)
  
      dataset = load_dataset("mnist")
      dataset.set_format("np")
      X_train = np.stack(dataset["train"]["image"])
      X_test = np.stack(dataset["test"]["image"])
  
      print("X_train:", X_train.shape, X_train.dtype)
      print("X_test:", X_test.shape, X_test.dtype)
  
      model = eg.Model(
          module=MLP(n1=256, n2=64),
          loss=MeanSquaredError(),
          optimizer=optax.rmsprop(0.001),
          eager=eager,
      )
  
      model.summary(X_train[:64])
  
      # Notice we are not passing `y`
      history = model.fit(
          inputs=X_train,
          epochs=epochs,
          steps_per_epoch=steps_per_epoch,
          batch_size=batch_size,
          validation_data=(X_test,),
          shuffle=True,
          callbacks=[
              eg.callbacks.TensorBoard(logdir=logdir, update_freq=300)
          ],
      )
  
      eg.utils.plot_history(history)
  
      # get random samples
      idxs = np.random.randint(0, 10000, size=(5,))
      x_sample = X_test[idxs]
  
      # get predictions
      y_pred = model.predict(x=x_sample)
  
      # plot and save results
      with SummaryWriter(os.path.join(logdir, "val")) as tbwriter:
          figure = plt.figure(figsize=(12, 12))
          for i in range(5):
              plt.subplot(2, 5, i + 1)
              plt.imshow(x_sample[i], cmap="gray")
              plt.subplot(2, 5, 5 + i + 1)
              plt.imshow(y_pred[i], cmap="gray")
  
      plt.show()
  
file_path: elegy/examples/elegy/mnist_autoencoder.py
imports: import os
  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 for calling a global function with example usage: main(debug, eager, logdir, steps_per_epoch, epochs, batch_size) without return types