class_docstr: Contains runtime information for any Determined workflow that uses the ``PyTorch`` API.
  
  With this class, users can do the following things:
  
  1. Wrap PyTorch models, optimizers, and LR schedulers with their Determined-compatible
     counterparts using :meth:`wrap_model`, :meth:`wrap_optimizer`, :meth:`wrap_lr_scheduler`,
     respectively. The Determined-compatible objects are capable of transparent
     distributed training, checkpointing and exporting, mixed-precision training,
     and gradient aggregation.
  2. Configure apex amp by calling :meth:`configure_apex_amp` (optional).
  3. Calculate the gradients with :meth:`backward` on a specified loss.
  4. Run an optimization step with :meth:`step_optimizer`.
  5. Functionalities inherited from :class:`determined.TrialContext`, including getting
     the runtime information and properly handling training data in distributed training.
class_docstr_tok: ["Contains","runtime","information","for","any","Determined","workflow","that","uses","the","``","PyTorch","``","API",".","With","this","class",",","users","can","do","the","following","things",":","1",".","Wrap","PyTorch","models",",","optimizers",",","and","LR","schedulers","with","their","Determined-compatible","counterparts","using",":","meth",":","`","wrap_model","`",",",":","meth",":","`","wrap_optimizer","`",",",":","meth",":","`","wrap_lr_scheduler","`",",","respectively",".","The","Determined-compatible","objects","are","capable","of","transparent","distributed","training",",","checkpointing","and","exporting",",","mixed-precision","training",",","and","gradient","aggregation",".","2",".","Configure","apex","amp","by","calling",":","meth",":","`","configure_apex_amp","`","(","optional",")",".","3",".","Calculate","the","gradients","with",":","meth",":","`","backward","`","on","a","specified","loss",".","4",".","Run","an","optimization","step","with",":","meth",":","`","step_optimizer","`",".","5",".","Functionalities","inherited","from",":","class",":","`","determined.TrialContext","`",",","including","getting","the","runtime","information","and","properly","handling","training","data","in","distributed","training","."]
docstr: Returns a wrapped optimizer.
  
  The optimizer must use the models wrapped by :meth:`wrap_model`. This function
  creates a ``horovod.DistributedOptimizer`` if using parallel/distributed training.
  
  ``backward_passes_per_step`` can be used to specify how many gradient aggregation
  steps will be performed in a single ``train_batch`` call per optimizer step.
  In most cases, this will just be the default value 1.  However, this advanced functionality
  can be used to support training loops like the one shown below:
  
  .. code-block:: python
  
      def train_batch(
          self, batch: TorchData, epoch_idx: int, batch_idx: int
      ) -> Dict[str, torch.Tensor]:
          data, labels = batch
          output = self.model(data)
          loss1 = output['loss1']
          loss2 = output['loss2']
          self.context.backward(loss1)
          self.context.backward(loss2)
          self.context.step_optimizer(self.optimizer, backward_passes_per_step=2)
          return {"loss1": loss1, "loss2": loss2}
code: def wrap_optimizer(
      self,
      optimizer: torch.optim.Optimizer,
      backward_passes_per_step: int = 1,
      fp16_compression: Optional[bool] = None,
      average_aggregated_gradients: Optional[bool] = None,
  ) -> torch.optim.Optimizer:
      """Returns a wrapped optimizer.
  
      The optimizer must use the models wrapped by :meth:`wrap_model`. This function
      creates a ``horovod.DistributedOptimizer`` if using parallel/distributed training.
  
      ``backward_passes_per_step`` can be used to specify how many gradient aggregation
      steps will be performed in a single ``train_batch`` call per optimizer step.
      In most cases, this will just be the default value 1.  However, this advanced functionality
      can be used to support training loops like the one shown below:
  
      .. code-block:: python
  
          def train_batch(
              self, batch: TorchData, epoch_idx: int, batch_idx: int
          ) -> Dict[str, torch.Tensor]:
              data, labels = batch
              output = self.model(data)
              loss1 = output['loss1']
              loss2 = output['loss2']
              self.context.backward(loss1)
              self.context.backward(loss2)
              self.context.step_optimizer(self.optimizer, backward_passes_per_step=2)
              return {"loss1": loss1, "loss2": loss2}
  
      """
      if self._managed_training:
          if self._use_apex:
              raise det.errors.InvalidExperimentException(
                  "Must call wrap_optimizer() before configure_apex_amp.",
              )
          if backward_passes_per_step < 1:
              raise det.errors.InvalidExperimentException(
                  "backward_passes_per_step for local gradient aggregation must be >= 1; "
                  f"got {backward_passes_per_step}.",
              )
  
          if (
              self.distributed.size > 1
              and self._distributed_backend.use_horovod()
          ):
              # We always override default fp16_compression setting if passed in directly
              if fp16_compression is None:
                  fp16_compression = self._fp16_compression_default
  
              optimizer = hvd.DistributedOptimizer(
                  optimizer,
                  named_parameters=self._filter_named_parameters(
                      optimizer
                  ),
                  backward_passes_per_step=backward_passes_per_step
                  * self._aggregation_frequency,
                  compression=hvd.Compression.fp16
                  if fp16_compression
                  else hvd.Compression.none,
              )
              logger.debug(
                  "Initialized optimizer for distributed and optimized parallel training."
              )
  
      if average_aggregated_gradients is None:
          average_aggregated_gradients = (
              self._average_aggregated_gradients_default
          )
  
      self._optimizer_configs[optimizer] = {
          "average_aggregated_gradients": average_aggregated_gradients
      }
      self.optimizers.append(optimizer)
      return optimizer
  
imports: import contextlib
  import logging
  import pathlib
  import warnings
  from typing import Any, Callable, Dict, Iterator, List, Optional, Set, Tuple, Type, Union
  import torch
  from torch import nn
  import determined
  from determined import profiler, pytorch, util
  from determined.horovod import hvd
label_desc: Use image node_id 13 for calling the PyTorchTrialContext obj's underlying member method code with example usage: obj.wrap_optimizer(optimizer, backward_passes_per_step, fp16_compression, average_aggregated_gradients) and returns: optimizer