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: Return parameter's per sample gradients as a single tensor.
  
  By default, per sample gradients (``p.grad_sample``) are stored as one tensor per
  batch basis. Therefore, ``p.grad_sample`` is a single tensor if holds results from
  only one batch, and a list of tensors if gradients are accumulated over multiple
  steps. This is done to provide visibility into which sample belongs to which batch,
  and how many batches have been processed.
  
  This method returns per sample gradients as a single concatenated tensor, regardless
  of how many batches have been accumulated
  
  Args:
      p: Parameter tensor. Must have ``grad_sample`` attribute
  
  Returns:
      ``p.grad_sample`` if it's a tensor already, or a single tensor computed by
      concatenating every tensor in ``p.grad_sample`` if it's a list
  
  Raises:
      ValueError
          If ``p`` is missing ``grad_sample`` attribute
docstr_tok: ["Return","parameter","'s","per","sample","gradients","as","a","single","tensor",".","By","default",",","per","sample","gradients","(","``","p.grad_sample","``",")","are","stored","as","one","tensor","per","batch","basis",".","Therefore",",","``","p.grad_sample","``","is","a","single","tensor","if","holds","results","from","only","one","batch",",","and","a","list","of","tensors","if","gradients","are","accumulated","over","multiple","steps",".","This","is","done","to","provide","visibility","into","which","sample","belongs","to","which","batch",",","and","how","many","batches","have","been","processed",".","This","method","returns","per","sample","gradients","as","a","single","concatenated","tensor",",","regardless","of","how","many","batches","have","been","accumulated","Args",":","p",":","Parameter","tensor",".","Must","have","``","grad_sample","``","attribute","Returns",":","``","p.grad_sample","``","if","it","'s","a","tensor","already",",","or","a","single","tensor","computed","by","concatenating","every","tensor","in","``","p.grad_sample","``","if","it","'s","a","list","Raises",":","ValueError","If","``","p","``","is","missing","``","grad_sample","``","attribute"]
code: def _get_flat_grad_sample(self, p: torch.Tensor):
      """
      Return parameter's per sample gradients as a single tensor.
  
      By default, per sample gradients (``p.grad_sample``) are stored as one tensor per
      batch basis. Therefore, ``p.grad_sample`` is a single tensor if holds results from
      only one batch, and a list of tensors if gradients are accumulated over multiple
      steps. This is done to provide visibility into which sample belongs to which batch,
      and how many batches have been processed.
  
      This method returns per sample gradients as a single concatenated tensor, regardless
      of how many batches have been accumulated
  
      Args:
          p: Parameter tensor. Must have ``grad_sample`` attribute
  
      Returns:
          ``p.grad_sample`` if it's a tensor already, or a single tensor computed by
          concatenating every tensor in ``p.grad_sample`` if it's a list
  
      Raises:
          ValueError
              If ``p`` is missing ``grad_sample`` attribute
      """
  
      if not hasattr(p, "grad_sample"):
          raise ValueError(
              "Per sample gradient not found. Are you using GradSampleModule?"
          )
      if p.grad_sample is None:
          raise ValueError(
              "Per sample gradient is not initialized. Not updated in backward pass?"
          )
      if isinstance(p.grad_sample, torch.Tensor):
          ret = p.grad_sample
      elif isinstance(p.grad_sample, list):
          ret = torch.cat(p.grad_sample, dim=0)
      else:
          raise ValueError(
              f"Unexpected grad_sample type: {type(p.grad_sample)}"
          )
  
      return ret
  
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 2 for calling the DPOptimizer obj's underlying member method code with example usage: obj._get_flat_grad_sample(p) and returns: ret