code: def main():
      parser = argparse.ArgumentParser(
          "Train a neural network on MNIST dataset"
      )
      parser.add_argument(
          "--batchsize", "-B", type=int, default=100, help="Batch size"
      )
      parser.add_argument(
          "--epoch",
          "-E",
          type=int,
          default=20,
          help="Number of epochs to train",
      )
      parser.add_argument(
          "--iteration",
          "-I",
          type=int,
          default=None,
          help="Number of iterations to train. Epoch is ignored if specified.",
      )
      parser.add_argument(
          "--data",
          "-p",
          default="mnist",
          help="Path to the directory that contains MNIST dataset",
      )
      parser.add_argument(
          "--device", "-d", default="native", help="Device to use"
      )
      parser.add_argument(
          "--eval-size",
          default=None,
          type=int,
          help="Number of samples to use from the test set for evaluation. "
          "None to use all.",
      )
      args = parser.parse_args()
  
      chx.set_default_device(args.device)
  
      # Prepare dataset
      X, Y = get_mnist(args.data, "train")
      X_test, Y_test = get_mnist(args.data, "t10k")
  
      # Prepare model
      model = MLP()
  
      # Training
      N = X.shape[0]  # TODO(beam2d): implement len
      # TODO(beam2d): support int32 indexing
      all_indices_np = np.arange(N, dtype=np.int64)
      batch_size = args.batchsize
      eval_size = args.eval_size
  
      # Train
      model.require_grad()
  
      it = 0
      epoch = 0
      is_finished = False
      start = time.time()
  
      while not is_finished:
          # TODO(beam2d): not suupported in chx
          np.random.shuffle(all_indices_np)
          all_indices = chx.array(all_indices_np)
  
          for i in range(0, N, batch_size):
              indices = all_indices[i : i + batch_size]
              x = X.take(indices, axis=0)
              t = Y.take(indices, axis=0)
  
              y = model.forward(x)
              loss = chx.softmax_cross_entropy(y, t).mean()
  
              loss.backward()
              model.update(lr=0.01)
  
              it += 1
              if args.iteration is not None:
                  mean_loss, accuracy = evaluate(
                      model, X_test, Y_test, eval_size, batch_size
                  )
                  elapsed_time = time.time() - start
                  print(
                      "iteration {}... loss={},\taccuracy={},\telapsed_time={}".format(
                          it, mean_loss, accuracy, elapsed_time
                      )
                  )
                  if it >= args.iteration:
                      is_finished = True
                      break
  
          epoch += 1
          if (
              args.iteration is None
          ):  # stop based on epoch, instead of iteration
              mean_loss, accuracy = evaluate(
                  model, X_test, Y_test, eval_size, batch_size
              )
              elapsed_time = time.time() - start
              print(
                  "epoch {}... loss={},\taccuracy={},\telapsed_time={}".format(
                      epoch, mean_loss, accuracy, elapsed_time
                  )
              )
              if epoch >= args.epoch:
                  is_finished = True
  
file_path: chainer/chainerx_cc/examples/mnist_py/train_mnist.py
imports: import argparse
  import gzip
  import pathlib
  import time
  import numpy
  import chainerx
label_desc: Use image node_id 3 for calling a global function with example usage: main() without return types