class_docstr: This class tests the crypten.nn package.
class_docstr_tok: ["This","class","tests","the","crypten.nn","package","."]
docstr: Tests all non-container Modules in crypten.nn that have equivalent
  modules in PyTorch.
code: def test_pytorch_modules(self):
      """
      Tests all non-container Modules in crypten.nn that have equivalent
      modules in PyTorch.
      """
  
      # input arguments for modules and input sizes:
      module_args = {
          "AdaptiveAvgPool2d": ((8, 8),),
          "AdaptiveMaxPool2d": ((2, 2),),
          "AvgPool2d": (2,),
          "BatchNorm1d": (25,),
          "BatchNorm2d": (3,),
          "BatchNorm3d": (6,),
          # "ConstantPad1d": (3, 1.0),
          # "ConstantPad2d": (2, 2.0),
          # "ConstantPad3d": (1, 0.0),   # TODO: Support negative steps in Slice.
          "Conv1d": (3, 6, 5),
          "Conv2d": (3, 6, 5),
          "Hardtanh": (-3, 1),
          "Linear": (400, 120),
          "MaxPool2d": (2,),
          "ReLU": (),
          "ReLU6": (),
          "Sigmoid": (),
          "Softmax": (0,),
          "LogSoftmax": (0,),
      }
      input_sizes = {
          "AdaptiveAvgPool2d": (1, 2, 24, 24),
          "AdaptiveMaxPool2d": (1, 3, 8, 8),
          "AvgPool2d": (1, 3, 32, 32),
          "BatchNorm1d": (8, 25),
          "BatchNorm2d": (8, 3, 7, 9),
          "BatchNorm3d": (8, 6, 3, 4, 2),
          "ConstantPad1d": (9,),
          "ConstantPad2d": (3, 6),
          "ConstantPad3d": (4, 2, 7),
          "Conv1d": (1, 3, 32),
          "Conv2d": (1, 3, 32, 32),
          "Hardtanh": (1, 3, 32, 32),
          "Linear": (1, 400),
          "LogSoftmax": (5, 5, 5),
          "MaxPool2d": (1, 2, 32, 32),
          "ReLU": (1, 3, 32, 32),
          "ReLU6": (1, 3, 32, 32),
          "Sigmoid": (8, 3, 32, 32),
          "Softmax": (5, 5, 5),
      }
      not_produced_by_onnx = [
          "BatchNorm1d",
          "BatchNorm2d",
          "BatchNorm3d",
          "Conv1d",
          "Conv2d",
          "Linear",
      ]
  
      # loop over all modules:
      for (
          module_name,
          from_pytorch,
          compute_gradients,
      ) in itertools.product(
          module_args.keys(), [False, True], [False, True]
      ):
          # some modules cannot be produced by the ONNX exporter:
          if from_pytorch and module_name in not_produced_by_onnx:
              continue
  
          # generate inputs:
          input = get_random_test_tensor(
              size=input_sizes[module_name], is_float=True
          )
          input.requires_grad = True
          encr_input = crypten.cryptensor(input)
          encr_input.requires_grad = compute_gradients
  
          # create PyTorch module:
          args = module_args[module_name]
          kwargs = (
              {"ceil_mode": True} if module_name == "MaxPool2d" else {}
          )
          module = getattr(torch.nn, module_name)(*args, **kwargs)
          module.train()
  
          # create encrypted CrypTen module:
          if from_pytorch:
              encr_module = crypten.nn.from_pytorch(module, input)
          else:
              encr_module = getattr(crypten.nn, module_name)(
                  *args, **kwargs
              )
              for name, param in module.named_parameters():
                  setattr(encr_module, name, param)
  
          # check that module properly encrypts / decrypts and
          # check that encrypting with current mode properly performs no-op
          for encrypted in [False, True, True, False, True]:
              encr_module.encrypt(mode=encrypted)
              if encrypted:
                  self.assertTrue(
                      encr_module.encrypted, "module not encrypted"
                  )
              else:
                  self.assertFalse(
                      encr_module.encrypted, "module encrypted"
                  )
  
              # check value of parameters:
              for key in ["weight", "bias"]:
                  if hasattr(module, key):  # if PyTorch model has key
                      # find that key in the crypten.nn.Graph:
                      if from_pytorch:
                          for encr_node in encr_module.modules():
                              if hasattr(encr_node, key):
                                  encr_param = getattr(encr_node, key)
                                  break
  
                      # or get it from the crypten Module directly:
                      else:
                          encr_param = getattr(encr_module, key)
  
                      # compare with reference:
                      # NOTE: Because some parameters are initialized randomly
                      # with different values on each process, we only want to
                      # check that they are consistent with source parameter value
                      reference = getattr(module, key)
                      src_reference = comm.get().broadcast(reference, 0)
                      msg = "parameter %s in %s incorrect" % (
                          key,
                          module_name,
                      )
                      if isinstance(encr_param, crypten.nn.Parameter):
                          encr_param = encr_param.data
                      if not crypten.is_encrypted_tensor(encr_param):
                          encr_param = crypten.cryptensor(
                              encr_param, src=0
                          )
                      self._check(encr_param, src_reference, msg)
  
          # Forward Pass
          self.assertTrue(
              encr_module.training, "training value incorrect"
          )
          reference = module(input)
          encr_output = encr_module(encr_input)
  
          msg = "from_pytorch" if from_pytorch else ""
          self._check(
              encr_output,
              reference,
              f"{module_name} forward failed {msg}",
          )
  
          # Backward Pass
          reference.sum().backward()
          encr_output.sum().backward()
  
          # Check input gradients
          if compute_gradients:
              self.assertIsNotNone(
                  encr_input.grad,
                  f"{module_name} grad failed to populate {msg}.",
              )
              self._check(
                  encr_input.grad,
                  input.grad,
                  f"{module_name} backward on input failed {msg}",
              )
          else:
              self.assertIsNone(encr_input.grad)
  
          # Check parameter gradients
          for name, encr_param in encr_module.named_parameters():
              name = name.split(".")[-1]
              torch_param = getattr(module, name)
              self._check(
                  encr_param.grad,
                  torch_param.grad,
                  f"{module_name} backward on parameter {name} failed {msg}",
              )
  
imports: import copy
  import itertools
  import logging
  import unittest
  import crypten
  import crypten.communicator
  import torch
  import torch.nn.functional
  from crypten.common.rng import generate_random_ring_element
  from crypten.common.tensor_types import is_float_tensor
  from crypten.config import cfg
  from crypten.encoder import FixedPointEncoder
  from test.multiprocess_test_case import get_random_linear, get_random_test_tensor, MultiProcessTestCase, onehot
label_desc: Use image node_id 9 for calling the TestNN obj's underlying member method code with example usage: obj.test_pytorch_modules() without return types