class_docstr: ``torch.optim.Optimizer`` wrapper that adds additional functionality to clip per
  sample gradients and add Gaussian noise.
  
  Can be used with any ``torch.optim.Optimizer`` subclass as an underlying optimizer.
  ``DPOptimzer`` assumes that parameters over which it performs optimization belong
  to GradSampleModule and therefore have the ``grad_sample`` attribute.
  
  On a high level ``DPOptimizer``'s step looks like this:
  1) Aggregate ``p.grad_sample`` over all parameters to calculate per sample norms
  2) Clip ``p.grad_sample`` so that per sample norm is not above threshold
  3) Aggregate clipped per sample gradients into ``p.grad``
  4) Add Gaussian noise to ``p.grad`` calibrated to a given noise multiplier and
  max grad norm limit (``std = noise_multiplier * max_grad_norm``).
  5) Call underlying optimizer to perform optimization step
  
  Examples:
      >>> module = MyCustomModel()
      >>> optimizer = torch.optim.SGD(module.parameters(), lr=0.1)
      >>> dp_optimizer = DPOptimizer(
      ...     optimizer=optimizer,
      ...     noise_multiplier=1.0,
      ...     max_grad_norm=1.0,
      ...     expected_batch_size=4,
      ... )
docstr: Clear gradients.
  
  Clears ``p.grad``, ``p.grad_sample`` and ``p.summed_grad`` for all of it's parameters
  
  Notes:
      ``set_to_none`` argument only affects ``p.grad``. ``p.grad_sample`` and
      ``p.summed_grad`` is never zeroed out and always set to None.
      Normal grads can do this, because their shape is always the same.
      Grad samples do not behave like this, as we accumulate gradients from different
      batches in a list
  
  Args:
      set_to_none: instead of setting to zero, set the grads to None. (only
      affects regular gradients. Per sample gradients are always set to None)
docstr_tok: ["Clear","gradients",".","Clears","``","p.grad","``",",","``","p.grad_sample","``","and","``","p.summed_grad","``","for","all","of","it","'s","parameters","Notes",":","``","set_to_none","``","argument","only","affects","``","p.grad","``",".","``","p.grad_sample","``","and","``","p.summed_grad","``","is","never","zeroed","out","and","always","set","to","None",".","Normal","grads","can","do","this",",","because","their","shape","is","always","the","same",".","Grad","samples","do","not","behave","like","this",",","as","we","accumulate","gradients","from","different","batches","in","a","list","Args",":","set_to_none",":","instead","of","setting","to","zero",",","set","the","grads","to","None",".","(","only","affects","regular","gradients",".","Per","sample","gradients","are","always","set","to","None",")"]
code: def zero_grad(self, set_to_none: bool = False):
      """
      Clear gradients.
  
      Clears ``p.grad``, ``p.grad_sample`` and ``p.summed_grad`` for all of it's parameters
  
      Notes:
          ``set_to_none`` argument only affects ``p.grad``. ``p.grad_sample`` and
          ``p.summed_grad`` is never zeroed out and always set to None.
          Normal grads can do this, because their shape is always the same.
          Grad samples do not behave like this, as we accumulate gradients from different
          batches in a list
  
      Args:
          set_to_none: instead of setting to zero, set the grads to None. (only
          affects regular gradients. Per sample gradients are always set to None)
      """
  
      if set_to_none is False:
          logger.debug(
              "Despite set_to_none is set to False, "
              "opacus will set p.grad_sample and p.summed_grad to None due to "
              "non-trivial gradient accumulation behaviour"
          )
  
      for p in self.params:
          p.grad_sample = None
  
          if not self._is_last_step_skipped:
              p.summed_grad = None
  
      self.original_optimizer.zero_grad(set_to_none)
  
imports: from __future__ import annotations
  import logging
  from typing import Callable, List, Optional, Union
  import torch
  from opacus.optimizers.utils import params
  from opt_einsum.contract import contract
  from torch import nn
  from torch.optim import Optimizer
label_desc: Use image node_id 12 for calling the DPOptimizer obj's underlying member method code with example usage: obj.zero_grad(set_to_none) without return types