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: Compute the gradient of current tensor w.r.t. graph leaves.
  
  The arguments are used in the same way as ``torch.Tensor.backward``.
  See https://pytorch.org/docs/1.4.0/_modules/torch/tensor.html#Tensor.backward for details.
  
  .. warning::
      When using distributed training, we don't support manual gradient accumulation.
      That means the gradient on each parameter can only be calculated once on each batch.
      If a parameter is associated with multiple losses, you can either choose to call
      ``backward'' on only one of those losses, or you can set the ``require_grads`` flag of
      a parameter or module to ``False`` to avoid manual gradient accumulation on that
      parameter.
      However, you can do gradient accumulation across batches by setting
      :ref:`optimizations.aggregation_frequency <config-aggregation-frequency>` in the
      experiment configuration to be greater than 1.
  
  Arguments:
      gradient (Tensor or None): Gradient w.r.t. the
          tensor. If it is a tensor, it will be automatically converted
          to a Tensor that does not require grad unless ``create_graph`` is True.
          None values can be specified for scalar Tensors or ones that
          don't require grad. If a None value would be acceptable then
          this argument is optional.
      retain_graph (bool, optional): If ``False``, the graph used to compute
          the grads will be freed. Note that in nearly all cases setting
          this option to True is not needed and often can be worked around
          in a much more efficient way. Defaults to the value of
          ``create_graph``.
      create_graph (bool, optional): If ``True``, graph of the derivative will
          be constructed, allowing to compute higher order derivative
          products. Defaults to ``False``.
code: def backward(
      self,
      loss: torch.Tensor,
      gradient: Optional[torch.Tensor] = None,
      retain_graph: bool = False,
      create_graph: bool = False,
  ) -> None:
      """Compute the gradient of current tensor w.r.t. graph leaves.
  
      The arguments are used in the same way as ``torch.Tensor.backward``.
      See https://pytorch.org/docs/1.4.0/_modules/torch/tensor.html#Tensor.backward for details.
  
      .. warning::
          When using distributed training, we don't support manual gradient accumulation.
          That means the gradient on each parameter can only be calculated once on each batch.
          If a parameter is associated with multiple losses, you can either choose to call
          ``backward'' on only one of those losses, or you can set the ``require_grads`` flag of
          a parameter or module to ``False`` to avoid manual gradient accumulation on that
          parameter.
          However, you can do gradient accumulation across batches by setting
          :ref:`optimizations.aggregation_frequency <config-aggregation-frequency>` in the
          experiment configuration to be greater than 1.
  
      Arguments:
          gradient (Tensor or None): Gradient w.r.t. the
              tensor. If it is a tensor, it will be automatically converted
              to a Tensor that does not require grad unless ``create_graph`` is True.
              None values can be specified for scalar Tensors or ones that
              don't require grad. If a None value would be acceptable then
              this argument is optional.
          retain_graph (bool, optional): If ``False``, the graph used to compute
              the grads will be freed. Note that in nearly all cases setting
              this option to True is not needed and often can be worked around
              in a much more efficient way. Defaults to the value of
              ``create_graph``.
          create_graph (bool, optional): If ``True``, graph of the derivative will
              be constructed, allowing to compute higher order derivative
              products. Defaults to ``False``.
      """
      if self._use_apex:
          if (
              self._last_backward_batch_idx is None
              or self._current_batch_idx is None
              or self._last_backward_batch_idx < self._current_batch_idx
          ):
              self._last_backward_batch_idx = self._current_batch_idx
          else:
              raise det.errors.InvalidExperimentException(
                  "Calling context.backward(loss) multiple times is not supported "
                  "while using apex.amp and parallel/distributed training"
              )
  
          if loss not in self._loss_ids:
              self._loss_ids[loss] = len(self._loss_ids)
          with apex.amp.scale_loss(
              loss, self.optimizers, loss_id=self._loss_ids[loss]
          ) as scaled_loss:
              with self._record_timing(
                  "train_batch.backward", accumulate=True
              ):
                  scaled_loss.backward(
                      gradient=gradient,
                      retain_graph=retain_graph,
                      create_graph=create_graph,
                  )
  
              if (
                  self.distributed.size > 1
                  and self._should_communicate_and_update()
                  and self._distributed_backend.use_horovod()
              ):
                  # When we exit out of this context manager, we need to finish
                  # communicating gradient updates before they are unscaled.
                  #
                  # Unfortunately, there is no clean way to support unscaling
                  # happening after synchronizing the optimizer on apex.amp.
                  # A short-term solution is to not support multiple losses nor
                  # multiple backward passes on one loss. A long-term solution is
                  # to integrate torch native AMP (https://pytorch.org/docs/stable/amp.html),
                  # which will come out soon.
                  with self._record_timing(
                      "train_batch.sync_optimizers", accumulate=True
                  ):
                      for optimizer in self.optimizers:
                          optimizer.synchronize()  # type: ignore
      else:
          if self._scaler and self.experimental._auto_amp:
              loss = self._scaler.scale(loss)
  
          with self._record_timing(
              "train_batch.backward", accumulate=True
          ):
              if (
                  self.distributed.size > 1
                  and self._distributed_backend.use_torch()
                  and not self._should_communicate_and_update()
              ):
                  # PyTorch DDP automatically syncs gradients by default on every backward pass.
                  # no_sync() disables gradient all-reduce until the last iteration.
                  with self._no_sync():
                      loss.backward(  # type: ignore
                          gradient=gradient,
                          retain_graph=retain_graph,
                          create_graph=create_graph,
                      )
              else:
                  loss.backward(  # type: ignore
                      gradient=gradient,
                      retain_graph=retain_graph,
                      create_graph=create_graph,
                  )
  
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 28 for calling the PyTorchTrialContext obj's underlying member method code with example usage: obj.backward(loss, gradient, retain_graph, create_graph) without return types