code: def main(
      steps_per_epoch: int = 200,
      batch_size: int = 64,
      epochs: int = 50,
      debug: bool = False,
      eager: bool = False,
      logdir: str = "runs",
  ):
      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")
      X_train = np.array(
          np.stack(dataset["train"]["image"]), dtype=np.uint8
      )
      X_test = np.array(
          np.stack(dataset["test"]["image"]), dtype=np.uint8
      )
      # Now binarize data
      X_train = (X_train > 0).astype(jnp.float32)
      X_test = (X_test > 0).astype(jnp.float32)
  
      print("X_train:", X_train.shape, X_train.dtype)
      print("X_test:", X_test.shape, X_test.dtype)
  
      model = eg.Model(
          module=VAE(latent_size=LATENT_SIZE),
          loss=[
              BinaryCrossEntropy(on="logits"),
              KL(weight=0.1),
          ],
          optimizer=optax.adam(1e-3),
          eager=eager,
      )
  
      model.summary(X_train[:batch_size])
  
      # Fit with datasets in memory
      history = model.fit(
          inputs=X_train,
          epochs=epochs,
          batch_size=batch_size,
          steps_per_epoch=steps_per_epoch,
          validation_data=(X_test,),
          shuffle=True,
          callbacks=[eg.callbacks.TensorBoard(logdir)],
      )
  
      print(
          "\n\n\nMetrics and images can be explored using tensorboard using:",
          f"\n \t\t\t tensorboard --logdir {logdir}",
      )
  
      eg.utils.plot_history(history)
  
      # get random samples
      idxs = np.random.randint(0, len(X_test), size=(5,))
      x_sample = X_test[idxs]
  
      # get predictions
      preds = model.predict(x=x_sample)
      y_pred = jax.nn.sigmoid(preds["logits"])
  
      # 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")
          # # tbwriter.add_figure("VAE Example", figure, epochs)
  
      plt.show()
  
imports: import os
  import typing
  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 import load_dataset
  from flax import linen
  from jax._src.numpy.lax_numpy import ndarray
  from tensorboardX.writer import SummaryWriter
  import elegy
label_desc: Use image node_id 1 for calling a global function with example usage: main(steps_per_epoch, batch_size, epochs, debug, eager, logdir) without return types