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 do not have
  equivalent modules in PyTorch.
code: def test_non_pytorch_modules(self):
      """
      Tests all non-container Modules in crypten.nn that do not have
      equivalent modules in PyTorch.
      """
  
      # input arguments for modules and input sizes:
      no_input_modules = ["Constant", "ConstantOfShape", "Range"]
      binary_modules = [
          "Add",
          "Concat",
          "Div",
          "Equal",
          "MatMul",
          "Mul",
          "Sub",
      ]
      ex_zero_modules = []
      module_args = {
          "Add": (),
          "Concat": (0,),
          "Constant": (1.2,),
          "ConstantOfShape": (1.4,),
          "Div": (),
          "Erf": (),
          "Equal": (),
          "Exp": (),
          "Expand": (),
          "Gather": (0,),
          "Gemm": (1.0, 1.0),
          "MatMul": (),
          "Mean": ([0], True),
          "Mul": (),
          "Pow": (),
          "Range": (),
          "Reshape": ((2, 2),),
          "Shape": (),
          "Slice": ([1], [4]),
          "Sqrt": (),
          "Sub": (),
          "Sum": ([0], True),
          "Squeeze": (0,),
          "Transpose": ([1, 3, 0, 2],),
          "Unsqueeze": (0,),
          "Where": (),
      }
      module_lambdas = {
          "Add": lambda x: x[0] + x[1],
          "Concat": lambda x: torch.cat((x[0], x[1])),
          "Constant": lambda _: torch.tensor(
              module_args["Constant"][0]
          ),
          "ConstantOfShape": lambda x: torch.tensor(
              module_args["ConstantOfShape"][0]
          ).expand(x[0]),
          "Div": lambda x: torch.div(x[0], x[1]),
          "Erf": lambda x: torch.erf(x),
          "Equal": lambda x: x[0].eq(x[1]),
          "Exp": lambda x: torch.exp(x),
          "Expand": lambda x: x[0].expand(x[1]),
          "Gather": lambda x: torch.from_numpy(
              x[0].numpy().take(x[1], module_args["Gather"][0])
          ),
          "Gemm": lambda x: x[0].matmul(x[1]).add(x[2]),
          "MatMul": lambda x: torch.matmul(x[0], x[1]),
          "Mean": lambda x: torch.mean(
              x,
              dim=module_args["Mean"][0],
              keepdim=(module_args["Mean"][1] == 1),
          ),
          "Mul": lambda x: x[0].mul(x[1]),
          "Pow": lambda x: x[0].pow(x[1]),
          "Range": lambda x: torch.arange(x[0], x[1], x[2]),
          "Reshape": lambda x: x[0].reshape(module_args["Reshape"][0]),
          "Shape": lambda x: torch.tensor(x.size()).float(),
          "Slice": lambda x: x[
              module_args["Slice"][0][0] : module_args["Slice"][1][0], :
          ],
          "Sqrt": lambda x: x.sqrt(),
          "Sub": lambda x: x[0] - x[1],
          "Sum": lambda x: torch.sum(
              x,
              dim=module_args["Sum"][0],
              keepdim=(module_args["Sum"][1] == 1),
          ),
          "Squeeze": lambda x: x.squeeze(module_args["Squeeze"][0]),
          "Transpose": lambda x: torch.from_numpy(
              x.numpy().transpose(module_args["Transpose"][0])
          ),
          "Unsqueeze": lambda x: x.unsqueeze(
              module_args["Unsqueeze"][0]
          ),
          "Where": lambda x: torch.where(x[0].byte(), x[1], x[2]),
      }
      additional_inputs = {
          "ConstantOfShape": ([2, 4],),
          "Expand": ([2, 4],),
          "Gather": (torch.tensor([[1, 2], [0, 3]]),),
          "Pow": (2,),
          "Range": (1, 6, 2),
      }
      input_sizes = {
          "Add": (10, 12),
          "Concat": (2, 2),
          "Constant": (1,),
          "Div": (3, 4),
          "Erf": (1, 2),
          "Equal": (2, 5, 3),
          "Exp": (10, 10, 10),
          "Expand": (1, 1),
          "Gather": (4, 4, 4, 4),
          "Gemm": (3, 4, 4),
          "MatMul": (2, 4, 4),
          "Mul": (4, 3, 2),
          "Mean": (3, 3, 3),
          "Pow": (4, 2),
          "Reshape": (1, 4),
          "Shape": (8, 3, 2),
          "Slice": (5, 2),
          "Sqrt": (2, 3),
          "Sub": (10, 12),
          "Sum": (3, 3, 3),
          "Squeeze": (1, 12, 6),
          "Transpose": (1, 2, 3, 4),
          "Unsqueeze": (8, 3),
          "Where": (3, 4, 2),
      }
      module_attributes = {
          # each attribute has two parameters: the name, and a bool indicating
          # whether the value should be wrapped into a list when the module is created
          "Concat": [("axis", False)],
          "Constant": [("value", False)],
          "ConstantOfShape": [("value", False)],
          "Gather": [("axis", False)],
          "Gemm": [("alpha", False), ("beta", False)],
          "Mean": [("axes", False), ("keepdims", False)],
          "Slice": [("starts", False), ("ends", False)],
          "Sum": [("axes", False), ("keepdims", False)],
          "Squeeze": [("axes", True)],
          "Transpose": [("perm", False)],
          "Unsqueeze": [("axes", True)],
      }
  
      # loop over all modules:
      for module_name in module_args.keys():
          # create encrypted CrypTen module:
          encr_module = getattr(crypten.nn, module_name)(
              *module_args[module_name]
          )
          encr_module.encrypt()
          self.assertTrue(encr_module.encrypted, "module not encrypted")
  
          # generate inputs:
          inputs, encr_inputs = None, None
          ex_zero_values = module_name in ex_zero_modules
          if module_name in binary_modules:
              inputs = [
                  get_random_test_tensor(
                      size=input_sizes[module_name],
                      is_float=True,
                      ex_zero=ex_zero_values,
                      max_value=1.0,
                  )
                  for _ in range(2)
              ]
              encr_inputs = [
                  crypten.cryptensor(input) for input in inputs
              ]
          elif module_name not in no_input_modules:
              inputs = get_random_test_tensor(
                  size=input_sizes[module_name],
                  is_float=True,
                  ex_zero=ex_zero_values,
                  max_value=1.0,
              )
              if (
                  module_name == "Where"
              ):  # Where condition is binary input
                  inputs = (inputs > 0.5).float()
              if module_name == "Sqrt":  # Sqrt requires positive input
                  inputs = inputs.abs()
              encr_inputs = crypten.cryptensor(inputs)
  
          # some modules take additional inputs:
          if module_name in additional_inputs:
              # base inputs:
              if inputs is None:
                  inputs, encr_inputs = [], []
              elif not isinstance(inputs, (list, tuple)):
                  inputs, encr_inputs = [inputs], [encr_inputs]
  
              # add additional inputs:
              for add_inp in additional_inputs[module_name]:
                  inputs.append(add_inp)
  
                  # encrypt only torch tensor inputs, not shapes or indices:
                  if torch.is_tensor(add_inp):
                      encr_inputs.append(crypten.cryptensor(add_inp))
                  else:
                      encr_inputs.append(add_inp)
  
          # some modules cannot work with encrypted inputs:
          if module_name in ["Gather"]:
              with self.assertRaises(ValueError):
                  encr_output = encr_module(encr_inputs)
  
              # but they can work using unencrypted indices:
              encr_inputs[1] = additional_inputs[module_name][0]
  
          # compare model outputs:
          reference = module_lambdas[module_name](inputs)
          encr_output = encr_module(encr_inputs)
          if torch.is_tensor(encr_output):
              self.assertTrue(
                  encr_module.SUPPORTS_PLAINTEXT_INPUTS,
                  msg=f"{module_name} has incorrect SUPPORTS_PLAINTEXT_INPUTS value",
              )
              encr_output = crypten.cryptensor(encr_output)
          self._check(encr_output, reference, "%s failed" % module_name)
  
          # create attributes for static from_onnx function
          local_attr = {}
          for i, attr_tuple in enumerate(
              module_attributes.get(module_name, [])
          ):
              attr_name, wrap_attr_in_list = attr_tuple
              if wrap_attr_in_list:
                  local_attr[attr_name] = [module_args[module_name][i]]
              else:
                  local_attr[attr_name] = module_args[module_name][i]
  
          # Update ReduceSum/ReduceMean module attributes, since the module and
          # from_onnx path are different
          if module_name == "ReduceSum":
              local_attr["keepdims"] = (
                  1 if module_args["ReduceSum"][1] is True else 0
              )
          if module_name == "ReduceMean":
              local_attr["keepdims"] = (
                  1 if module_args["ReduceMean"][1] is True else 0
              )
          if module_name == "Reshape":
              local_attr["shape"] = module_args["Reshape"][0]
  
          # compare model outputs using the from_onnx static function
          module = getattr(crypten.nn, module_name).from_onnx(
              attributes=local_attr
          )
          encr_module_onnx = module.encrypt()
          encr_output = encr_module_onnx(encr_inputs)
          if torch.is_tensor(encr_output):
              self.assertTrue(
                  encr_module_onnx.SUPPORTS_PLAINTEXT_INPUTS,
                  msg=f"{module_name} has incorrect SUPPORTS_PLAINTEXT_INPUTS value",
              )
              encr_output = crypten.cryptensor(encr_output)
          self._check(encr_output, reference, "%s failed" % module_name)
  
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 8 for calling the TestNN obj's underlying member method code with example usage: obj.test_non_pytorch_modules() without return types