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 and backward through both the standard nn_module and the dp_module and
  checks that all gradients 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``
      train_fn: A function that we can run to train the model on a single input batch.
          It should run forward and backward and stop there.
          Refer to ``default_train_fn`` in this file for an example.
      *train_fn_args: args to pass to the train_fn (ie we'll call
          ``train_fn(*train_fn_args, **train_fn_kwargs))``.
      atol: Absolute tolerance.
      rtol: Relative tolerance.
      **train_fn_kwargs: kwargs to pass to the train_fn (ie we'll call
          ``module(*module_args, **module_kwargs))``.
  Returns:
      None
  
  Raises:
      AssertionError if either:
          - nn_module has a gradient for a parameter that dp_module doesn't have.
          - dp_module has a gradient for a parameter that nn_module doesn't have.
          - The shape of any parameter gradient doesn't match.
          - The values for any parameter gradient ``nn_grad`` differ by more than
              `atol + rtol * abs(nn_grad)`.
docstr_tok: ["Runs","forward","and","backward","through","both","the","standard","nn_module","and","the","dp_module","and","checks","that","all","gradients","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","``","train_fn",":","A","function","that","we","can","run","to","train","the","model","on","a","single","input","batch",".","It","should","run","forward","and","backward","and","stop","there",".","Refer","to","``","default_train_fn","``","in","this","file","for","an","example",".","*","train_fn_args",":","args","to","pass","to","the","train_fn","(","ie","we","'ll","call","``","train_fn","(","*","train_fn_args",",","*","*","train_fn_kwargs",")",")","``",".","atol",":","Absolute","tolerance",".","rtol",":","Relative","tolerance",".","*","*","train_fn_kwargs",":","kwargs","to","pass","to","the","train_fn","(","ie","we","'ll","call","``","module","(","*","module_args",",","*","*","module_kwargs",")",")","``",".","Returns",":","None","Raises",":","AssertionError","if","either",":","-","nn_module","has","a","gradient","for","a","parameter","that","dp_module","does","n't","have",".","-","dp_module","has","a","gradient","for","a","parameter","that","nn_module","does","n't","have",".","-","The","shape","of","any","parameter","gradient","does","n't","match",".","-","The","values","for","any","parameter","gradient","``","nn_grad","``","differ","by","more","than","`","atol","+","rtol","*","abs","(","nn_grad",")","`","."]
code: def compare_gradients(
      self,
      nn_module: nn.Module,
      dp_module: nn.Module,
      train_fn: Callable,
      *train_fn_args,
      atol: float = 10e-6,
      rtol: float = 10e-5,
      **train_fn_kwargs,
  ) -> None:
      """
      Runs forward and backward through both the standard nn_module and the dp_module and
      checks that all gradients 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``
          train_fn: A function that we can run to train the model on a single input batch.
              It should run forward and backward and stop there.
              Refer to ``default_train_fn`` in this file for an example.
          *train_fn_args: args to pass to the train_fn (ie we'll call
              ``train_fn(*train_fn_args, **train_fn_kwargs))``.
          atol: Absolute tolerance.
          rtol: Relative tolerance.
          **train_fn_kwargs: kwargs to pass to the train_fn (ie we'll call
              ``module(*module_args, **module_kwargs))``.
      Returns:
          None
  
      Raises:
          AssertionError if either:
              - nn_module has a gradient for a parameter that dp_module doesn't have.
              - dp_module has a gradient for a parameter that nn_module doesn't have.
              - The shape of any parameter gradient doesn't match.
              - The values for any parameter gradient ``nn_grad`` differ by more than
                  `atol + rtol * abs(nn_grad)`.
      """
  
      train_fn(nn_module, *train_fn_args, **train_fn_kwargs)
      train_fn(dp_module, *train_fn_args, **train_fn_kwargs)
  
      nn_params = nn_module.state_dict()
      dp_params = dp_module.state_dict()
  
      nn_only_grads = [
          param_name
          for param_name, param in nn_params.items()
          if param.grad is not None and param_name not in dp_params
      ]
      if nn_only_grads:
          failed_str = "\n\t".join(
              f"{i}. {s}" for i, s in enumerate(nn_only_grads, 1)
          )
          raise AssertionError(
              f"A total of {len(nn_only_grads)} gradients are in nn_module "
              f"but not in dp_module: \n\t{failed_str}"
          )
  
      dp_only_grads = [
          param_name
          for param_name, param in dp_params.items()
          if param.grad is not None and param_name not in nn_params
      ]
      if dp_only_grads:
          failed_str = "\n\t".join(
              f"{i}. {s}" for i, s in enumerate(nn_only_grads, 1)
          )
          raise AssertionError(
              f"A total of {len(dp_only_grads)} gradients are in dp_module "
              f"but not in nn_module: \n\t{failed_str}"
          )
  
      for param_name, nn_param in nn_params.items():
          dp_param = dp_params[param_name]
          self._check_shapes((nn_param), (dp_param), (param_name))
          self._check_values(
              (nn_param), (dp_param), atol, rtol, (param_name)
          )
  
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 2 for calling the DPModules_test obj's underlying member method code with example usage: obj.compare_gradients(nn_module, dp_module, train_fn) without return types