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: Perform a single optimization step.
  
  This function must be called once for each optimizer. However, the order of
  different optimizers' steps can be specified by calling this function in different
  orders. Also, gradient accumulation across iterations is performed by the Determined
  training loop by setting the experiment configuration field
  :ref:`optimizations.aggregation_frequency <config-aggregation-frequency>`.
  
  Here is a code example:
  
  .. code-block:: python
  
      def clip_grads(params):
          torch.nn.utils.clip_grad_norm_(params, 0.0001),
  
      self.context.step_optimizer(self.opt1, clip_grads)
  
  Arguments:
      optimizer(``torch.optim.Optimizer``): Which optimizer should be stepped.
      clip_grads(a function, optional): This function should have one argument for
          parameters in order to clip the gradients.
      auto_zero_grads(bool, optional): Automatically zero out gradients automatically after
          stepping the optimizer. If false, you need to call ``optimizer.zero_grad()``
          manually. Note that if :ref:`optimizations.aggregation_frequency
          <config-aggregation-frequency>` is greater than 1, ``auto_zero_grads`` must be true.
      scaler(``torch.cuda.amp.GradScaler``, optional): The scaler to use for stepping the
          optimizer. This should be unset if not using AMP, and is necessary if
          ``wrap_scaler()`` was called directly.
code: def step_optimizer(
      self,
      optimizer: torch.optim.Optimizer,
      clip_grads: Optional[Callable[[Iterator], None]] = None,
      auto_zero_grads: bool = True,
      scaler: Optional[Any] = None,
      # Should be ``torch.cuda.amp.GradScaler``, but:
      #   * other implementations might be possible
      #   * requiring this type forces upgrades to PyTorch 1.6+
  ) -> None:
      """
      Perform a single optimization step.
  
      This function must be called once for each optimizer. However, the order of
      different optimizers' steps can be specified by calling this function in different
      orders. Also, gradient accumulation across iterations is performed by the Determined
      training loop by setting the experiment configuration field
      :ref:`optimizations.aggregation_frequency <config-aggregation-frequency>`.
  
      Here is a code example:
  
      .. code-block:: python
  
          def clip_grads(params):
              torch.nn.utils.clip_grad_norm_(params, 0.0001),
  
          self.context.step_optimizer(self.opt1, clip_grads)
  
      Arguments:
          optimizer(``torch.optim.Optimizer``): Which optimizer should be stepped.
          clip_grads(a function, optional): This function should have one argument for
              parameters in order to clip the gradients.
          auto_zero_grads(bool, optional): Automatically zero out gradients automatically after
              stepping the optimizer. If false, you need to call ``optimizer.zero_grad()``
              manually. Note that if :ref:`optimizations.aggregation_frequency
              <config-aggregation-frequency>` is greater than 1, ``auto_zero_grads`` must be true.
          scaler(``torch.cuda.amp.GradScaler``, optional): The scaler to use for stepping the
              optimizer. This should be unset if not using AMP, and is necessary if
              ``wrap_scaler()`` was called directly.
      """
  
      if self._aggregation_frequency > 1 and not auto_zero_grads:
          raise det.errors.InvalidExperimentException(
              "if optimizations.aggregation_frequency is larger than 1, "
              "auto_zero_grads must be set to true. ",
          )
  
      if not self._should_communicate_and_update():
          return
  
      # Communication needs to be synchronized so that is completed
      # before we apply gradient clipping and `step()`.
      # In the case of APEX this is called in backward() instead, so that it's inside the context
      # manager and before unscaling.
      # In the case of PyTorch DDP, losses are synchronized during the backwards() pass.
      if (
          self.distributed.size > 1
          and self._distributed_backend.use_horovod()
          and not self._use_apex
      ):
          with self._record_timing(
              "train_batch.sync_optimizers", accumulate=True
          ):
              optimizer.synchronize()  # type: ignore
  
      parameters = (
          [
              p
              for group in optimizer.param_groups
              for p in group.get("params", [])
          ]
          if not self._use_apex
          else apex.amp.master_params(optimizer)
      )
  
      if bool(
          self._optimizer_configs[optimizer][
              "average_aggregated_gradients"
          ]
      ):
          self._average_gradients(
              parameters=parameters, divisor=self._aggregation_frequency
          )
  
      if clip_grads is not None:
          if self._scaler and self.experimental._auto_amp:
              self._scaler.unscale_(optimizer)
          clip_grads(parameters)
  
      # For stepping the optimizer we will operate on the scaler passed
      # in, or fall back to the wrapped scaler (if any).
      if scaler is None and self.experimental._auto_amp:
          scaler = self._scaler
      if scaler:
  
          def step_fn() -> None:
              scaler.step(optimizer)  # type: ignore
  
      else:
          step_fn = optimizer.step  # type: ignore
  
      # In the case of PyTorch DDP, losses are synchronized automatically on the backwards() pass
      if (
          self.distributed.size > 1
          and self._distributed_backend.use_horovod()
      ):
          with optimizer.skip_synchronize():  # type: ignore
              step_fn()
      else:
          step_fn()
  
      if auto_zero_grads:
          optimizer.zero_grad()
  
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 30 for calling the PyTorchTrialContext obj's underlying member method code with example usage: obj.step_optimizer(optimizer, clip_grads, auto_zero_grads, scaler) without return types