def _run_one_step(self, task_index, task_changed):
    """Updates model weights/state and optimizer slots by running one step.

    Args:
      task_index (int): Index of the task to train on.
      task_changed (bool): Whether the state has changed since the last step.

    Returns:
      Tuple (loss, stats) with loss value from one step
      of training and stats, the current optimizer statistics.
    """
    step = self.step
    for callback in self._callbacks:
      if callback.call_at(step):
        callback.on_step_begin(step)

    learning_rate = self._tasks[task_index].learning_rate(step)
    batch = self._tasks[task_index].next_batch()
    rng = self.new_rng()
    trainer = self._trainer_per_task[task_index]
    if task_changed:
      # Re-replicate weights and state to synchronize them between tasks.
      self.update_weights_and_state(self._model.weights, self._model.state)

    (loss, stats) = trainer.one_step(
        batch, rng, step=step, learning_rate=learning_rate
    )

    for callback in self._callbacks:
      if callback.call_at(step):
        callback.on_step_end(step)

    return (loss, stats)
Updates model weights/state and optimizer slots by running one step.

    Args:
      task_index (int): Index of the task to train on.
      task_changed (bool): Whether the state has changed since the last step.

    Returns:
      Tuple (loss, stats) with loss value from one step
      of training and stats, the current optimizer statistics.
    
trax/supervised/training.py