class_docstr: Set of common testing utils. It is meant to be subclassed by your test.
  See other tests as an example of how this is done.
  
  The objective of these tests is to make sure that our DP-friendly reimplementations of
  standard nn.Modules such as LSTM are indeed drop-in replacements: we are checking that all
  outputs and states are the same between the two implementations. Here, we do NOT test for
  grad_samples, which is something we do in the grad_sample tests.
class_docstr_tok: ["Set","of","common","testing","utils",".","It","is","meant","to","be","subclassed","by","your","test",".","See","other","tests","as","an","example","of","how","this","is","done",".","The","objective","of","these","tests","is","to","make","sure","that","our","DP-friendly","reimplementations","of","standard","nn.Modules","such","as","LSTM","are","indeed","drop-in","replacements",":","we","are","checking","that","all","outputs","and","states","are","the","same","between","the","two","implementations",".","Here",",","we","do","NOT","test","for","grad_samples",",","which","is","something","we","do","in","the","grad_sample","tests","."]
docstr: Runs forward through both the standard nn_module and the dp_module and checks that all
  outputs are indeed the same.
  
  Args:
      nn_module: The original ``nn.Module`` that will be our reference
      dp_module: Our ``dp_module`` reimplementation that we want to test against ``nn_module``
      *module_args: args to pass to the model's forward (ie we'll call
          ``module(*module_args, **module_kwargs))``.
      output_names: If provided, will make printing more informative (rather than say
          output number 3 does not match" we can say "output `h` does not match").
      atol: Absolute tolerance.
      rtol: Relative tolerance.
      **module_kwargs: kwargs to pass to the model's forward (ie we'll call
          ``module(*module_args, **module_kwargs))``.
  Returns:
      None
  
  Raises:
      AssertionError if either:
          - The number of outputs of the forward doesn't match
          - The shape of any output doesn't match
          - The values for any output ``nn_out`` in ``nn_outs`` differ by more
              than `atol + rtol * abs(nn_out)`
docstr_tok: ["Runs","forward","through","both","the","standard","nn_module","and","the","dp_module","and","checks","that","all","outputs","are","indeed","the","same",".","Args",":","nn_module",":","The","original","``","nn.Module","``","that","will","be","our","reference","dp_module",":","Our","``","dp_module","``","reimplementation","that","we","want","to","test","against","``","nn_module","``","*","module_args",":","args","to","pass","to","the","model","'s","forward","(","ie","we","'ll","call","``","module","(","*","module_args",",","*","*","module_kwargs",")",")","``",".","output_names",":","If","provided",",","will","make","printing","more","informative","(","rather","than","say","output","number","3","does","not","match","''","we","can","say","``","output","`","h","`","does","not","match","''",")",".","atol",":","Absolute","tolerance",".","rtol",":","Relative","tolerance",".","*","*","module_kwargs",":","kwargs","to","pass","to","the","model","'s","forward","(","ie","we","'ll","call","``","module","(","*","module_args",",","*","*","module_kwargs",")",")","``",".","Returns",":","None","Raises",":","AssertionError","if","either",":","-","The","number","of","outputs","of","the","forward","does","n't","match","-","The","shape","of","any","output","does","n't","match","-","The","values","for","any","output","``","nn_out","``","in","``","nn_outs","``","differ","by","more","than","`","atol","+","rtol","*","abs","(","nn_out",")","`"]
code: def compare_forward_outputs(
      self,
      nn_module: nn.Module,
      dp_module: nn.Module,
      *module_args,
      output_names: Optional[Tuple[str, ...]] = None,
      atol: float = 10e-6,
      rtol: float = 10e-5,
      **module_kwargs,
  ) -> None:
      """
      Runs forward through both the standard nn_module and the dp_module and checks that all
      outputs are indeed the same.
  
      Args:
          nn_module: The original ``nn.Module`` that will be our reference
          dp_module: Our ``dp_module`` reimplementation that we want to test against ``nn_module``
          *module_args: args to pass to the model's forward (ie we'll call
              ``module(*module_args, **module_kwargs))``.
          output_names: If provided, will make printing more informative (rather than say
              output number 3 does not match" we can say "output `h` does not match").
          atol: Absolute tolerance.
          rtol: Relative tolerance.
          **module_kwargs: kwargs to pass to the model's forward (ie we'll call
              ``module(*module_args, **module_kwargs))``.
      Returns:
          None
  
      Raises:
          AssertionError if either:
              - The number of outputs of the forward doesn't match
              - The shape of any output doesn't match
              - The values for any output ``nn_out`` in ``nn_outs`` differ by more
                  than `atol + rtol * abs(nn_out)`
      """
      torch.use_deterministic_algorithms(True)
      torch.manual_seed(0)
      np.random.seed(0)
  
      batch_first_nn, batch_first_dp = (
          getattr(nn_module, "batch_first", None),
          getattr(dp_module, "batch_first", None),
      )
  
      nn_outs = nn_module(*module_args, **module_kwargs)
      nn_len = 0
      try:
          nn_len = len(nn_outs)
      except AttributeError:
          nn_outs = [nn_outs]
      nn_outs = flatten(nn_outs)
  
      dp_outs = dp_module(*module_args, **module_kwargs)
      dp_len = 0
      try:
          dp_len = len(dp_outs)
      except AttributeError:
          dp_outs = [dp_outs]
      dp_outs = flatten(dp_outs)
  
      self.assertEqual(
          dp_len,
          nn_len,
          f"The number of outputs does not match! Reference nn has {nn_len} outputs, and "
          f"DP reimplementation has {dp_len} outputs",
      )
  
      self._check_shapes(nn_outs, dp_outs, output_names=output_names)
      self._check_values(
          nn_outs,
          dp_outs,
          atol,
          rtol,
          output_names=output_names,
          batch_first_nn=batch_first_nn,
          batch_first_dp=batch_first_dp,
      )
  
file_path: opacus/opacus/tests/dp_layers/common.py
imports: import io
  import unittest
  from typing import Callable, Optional, Sequence, Tuple, Union
  import numpy
  import torch
  import torch.nn
  import torch.nn.functional
  from torch.nn.utils.rnn import PackedSequence, pad_packed_sequence
  from torch.testing import assert_close
label_desc: Use image node_id 1 for calling the DPModules_test obj's underlying member method code with example usage: obj.compare_forward_outputs(nn_module, dp_module) without return types