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"])
      y_train = dataset["train"]["label"]
      X_test = np.stack(dataset["test"]["image"])
      y_test = dataset["test"]["label"]
  
      print("X_train:", X_train.shape, X_train.dtype)
      print("y_train:", y_train.shape, y_train.dtype)
      print("X_test:", X_test.shape, X_test.dtype)
      print("y_test:", y_test.shape, y_test.dtype)
  
      model = Model(
          features_out=10,
          optimizer=optax.adam(1e-3),
          eager=eager,
      )
  
      history = model.fit(
          inputs=X_train,
          labels=y_train,
          epochs=epochs,
          steps_per_epoch=steps_per_epoch,
          batch_size=batch_size,
          validation_data=(X_test, y_test),
          shuffle=True,
          callbacks=[eg.callbacks.TensorBoard(logdir=logdir)],
      )
  
      eg.utils.plot_history(history)
  
file_path: elegy/examples/jax/linear_classifier_train_step.py
imports: import os
  import typing
  from datetime import datetime
  import jax
  import jax.numpy
  import numpy
  import optax
  import typer
  from datasets.load import load_dataset
  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