def __init__(self, labeled_data, loss_layer, optimizer,
               lr_schedule=None, n_steps_per_checkpoint=100,
               n_steps_per_permanent_checkpoint=None, loss_name=None,
               sample_batch=None, export_prefix=None):
    r"""Configures a training task.

    Args:
      labeled_data: Iterator of batches of labeled data tuples. Each tuple has
          1+ data (input value) tensors followed by 1 label (target value)
          tensor.  All tensors are NumPy ndarrays or their JAX counterparts.
      loss_layer: Layer that computes a scalar value (the "loss") by comparing
          model output :math:`\hat{y}=f(x)` to the target :math:`y`.
      optimizer: Optimizer object that computes model weight updates from
          loss-function gradients.
      lr_schedule: Learning rate schedule, a function step -> learning_rate.
      n_steps_per_checkpoint: How many steps to run between checkpoints.
      n_steps_per_permanent_checkpoint: How many steps to run between permanent
          checkpoints.
      loss_name: Name for the loss metric.
      sample_batch: Optional sample batch for model initialization. If not
          provided, it will be taken from ``labeled_data``.
      export_prefix: Optional task name to be used as prefix for exporting
      metrics during training in Loop.
    """
    self._export_prefix = export_prefix
    self._labeled_data = labeled_data
    self._loss_layer = loss_layer
    self._optimizer = optimizer
    self._lr_schedule = lr_schedule
    self._sample_batch = sample_batch or next(labeled_data)
    self._n_steps_per_checkpoint = n_steps_per_checkpoint
    self._n_steps_per_permanent_checkpoint = n_steps_per_permanent_checkpoint
    self._loss_name = loss_name or self._loss_layer.name
Configures a training task.

    Args:
      labeled_data: Iterator of batches of labeled data tuples. Each tuple has
          1+ data (input value) tensors followed by 1 label (target value)
          tensor.  All tensors are NumPy ndarrays or their JAX counterparts.
      loss_layer: Layer that computes a scalar value (the "loss") by comparing
          model output :math:`\hat{y}=f(x)` to the target :math:`y`.
      optimizer: Optimizer object that computes model weight updates from
          loss-function gradients.
      lr_schedule: Learning rate schedule, a function step -> learning_rate.
      n_steps_per_checkpoint: How many steps to run between checkpoints.
      n_steps_per_permanent_checkpoint: How many steps to run between permanent
          checkpoints.
      loss_name: Name for the loss metric.
      sample_batch: Optional sample batch for model initialization. If not
          provided, it will be taken from ``labeled_data``.
      export_prefix: Optional task name to be used as prefix for exporting
      metrics during training in Loop.
    
trax/supervised/training.py