text: # pylint: disable=unused-variable,arguments-differ
  
  # These should probably all live in separate files
  from typing import Set, Dict, TYPE_CHECKING
  import logging
  
  import torch
  
  from allennlp.common.params import Params
  from allennlp.training import util as training_util
  from allennlp.training.callbacks.callback import Callback, handle_event
  from allennlp.training.callbacks.events import Events
  from allennlp.training.tensorboard_writer import TensorboardWriter
  
  if TYPE_CHECKING:
      from allennlp.training.callback_trainer import CallbackTrainer  # pylint:disable=unused-import
  
  logger = logging.getLogger(__name__)
  
  
  @Callback.register("log_to_tensorboard")
  class LogToTensorboard(Callback):
      """
      Callback that handles all Tensorboard logging.
  
      Parameters
      ----------
      tensorboard : ``TensorboardWriter``
          The TensorboardWriter instance to write to.
      log_batch_size_period : int, optional (default: None)
          If provided, we'll log the average batch sizes to Tensorboard
          every this-many batches.
      """
      def __init__(self,
                   tensorboard: TensorboardWriter,
                   log_batch_size_period: int = None) -> None:
          self.log_batch_size_period = log_batch_size_period
          self.tensorboard = tensorboard
  
          self.cumulative_batch_size = 0
  
          # For logging histograms
          self.histogram_parameters: Set[str] = set()
          self.param_updates: Dict[str, torch.Tensor] = {}
  
      @handle_event(Events.TRAINING_START)
      def training_start(self, trainer: 'CallbackTrainer'):
          # pylint: disable=protected-access
  
          # This is an ugly hack to get the tensorboard instance to know about the trainer, because
          # the callbacks are defined before the trainer.
          # TODO: figure out a better way to handle this.
          self.tensorboard._get_batch_num_total = lambda: trainer.batch_num_total
  
          # Get histogram parameters
          self.histogram_parameters = set(
                  trainer.model.get_parameters_for_histogram_tensorboard_logging()
          )
  
          # Enable activation logging.
          if self.tensorboard._histogram_interval is not None:
              self.tensorboard.enable_activation_logging(trainer.model)
  
      @handle_event(Events.BATCH_START)
      def copy_current_parameters(self, trainer: 'CallbackTrainer'):
          if self.tensorboard.should_log_histograms_this_batch():
              # Get the magnitude of parameter updates for logging
              # We need a copy of current parameters to compute magnitude of updates,
              # and copy them to CPU so large models won't go OOM on the GPU.
              self.param_updates = {name: param.detach().cpu().clone()
                                    for name, param in trainer.model.named_parameters()}
  
      @handle_event(Events.BATCH_END)
      def batch_end_logging(self, trainer: 'CallbackTrainer'):
          # Log parameter values to tensorboard
          if self.tensorboard.should_log_this_batch():
              self.tensorboard.log_parameter_and_gradient_statistics(trainer.model, trainer.batch_grad_norm)
              self.tensorboard.log_learning_rates(trainer.model, trainer.optimizer)
  
              self.tensorboard.add_train_scalar("loss/loss_train", trainer.train_metrics["loss"])
              self.tensorboard.log_metrics({"epoch_metrics/" + k: v for k, v in trainer.train_metrics.items()})
  
  
          if self.log_batch_size_period:
              cur_batch = sum([training_util.get_batch_size(batch) for batch in trainer.batch_group])
              self.cumulative_batch_size += cur_batch
              if (trainer.batches_this_epoch - 1) % self.log_batch_size_period == 0:
                  average = self.cumulative_batch_size / trainer.batches_this_epoch
                  logger.debug(f"current batch size: {cur_batch} mean batch size: {average}")
                  self.tensorboard.add_train_scalar("current_batch_size", cur_batch)
                  self.tensorboard.add_train_scalar("mean_batch_size", average)
  
          if self.tensorboard.should_log_histograms_this_batch():
              for name, param in trainer.model.named_parameters():
                  self.param_updates[name].sub_(param.detach().cpu())
                  update_norm = torch.norm(self.param_updates[name].view(-1, ))
                  param_norm = torch.norm(param.view(-1, )).cpu()
                  self.tensorboard.add_train_scalar("gradient_update/" + name,
                                                    update_norm / (param_norm + 1e-7))
              self.param_updates.clear()
              self.tensorboard.log_histograms(trainer.model, self.histogram_parameters)
  
      @handle_event(Events.EPOCH_END)
      def epoch_end_logging(self, trainer: 'CallbackTrainer'):
          self.tensorboard.log_metrics(trainer.train_metrics,
                                       val_metrics=trainer.val_metrics,
                                       log_to_console=True,
                                       epoch=trainer.epoch_number + 1)
  
      @handle_event(Events.TRAINING_END)
      def training_end(self, trainer: 'CallbackTrainer'):
          # pylint: disable=unused-argument
          self.tensorboard.close()
  
      @classmethod
      def from_params(cls, serialization_dir: str, params: Params) -> 'LogToTensorboard':  # type: ignore
          log_batch_size_period = params.pop_int("log_batch_size_period", None)
          tensorboard = TensorboardWriter.from_params(params=params,
                                                      serialization_dir=serialization_dir,
                                                      get_batch_num_total=lambda: None)
          return LogToTensorboard(tensorboard, log_batch_size_period)
  