def _mnist_tasks(head=None):
  """Creates MNIST training and evaluation tasks.

  Args:
    head: Adaptor layer to put before loss and accuracy layers in the tasks.

  Returns:
    A pair (train_task, eval_task) consisting of the MNIST training task and the
    MNIST evaluation task using cross-entropy as loss and accuracy as metric.
  """
  loss = tl.WeightedCategoryCrossEntropy()
  accuracy = tl.WeightedCategoryAccuracy()
  if head is not None:
    loss = tl.Serial(head, loss)
    accuracy = tl.Serial(head, accuracy)
  task = training.TrainTask(
      itertools.cycle(_mnist_dataset().train_stream(1)),
      loss,
      adam.Adam(0.001),
  )
  eval_task = training.EvalTask(
      itertools.cycle(_mnist_dataset().eval_stream(1)),
      [loss, accuracy],
      n_eval_batches=10,
      metric_names=['CrossEntropy', 'WeightedCategoryAccuracy'],
  )
  return (task, eval_task)
Creates MNIST training and evaluation tasks.

  Args:
    head: Adaptor layer to put before loss and accuracy layers in the tasks.

  Returns:
    A pair (train_task, eval_task) consisting of the MNIST training task and the
    MNIST evaluation task using cross-entropy as loss and accuracy as metric.
  
trax/supervised/mnist_test.py