code: 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)
  
file_path: trax/trax/supervised/mnist_test.py
imports: import io
  import itertools
  from unittest import mock
  from absl.testing import absltest
  from trax import layers
  from trax.data import inputs
  from trax.data import tf_inputs
  from trax.optimizers import adam
  from trax.supervised import training
label_desc: Use image node_id 4 for calling a global function with example usage: _mnist_tasks(head) and returns: task,  eval_task