class_docstr: This class tests the crypten.nn package.
class_docstr_tok: ["This","class","tests","the","crypten.nn","package","."]
docstr: Tests dumping and loading of state dicts.
code: def test_state_dict(self) -> None:
      """
      Tests dumping and loading of state dicts.
      """
      import io
  
      def _check_equal(t1, t2):
          """
          Checks whether to tensors are identical.
          """
          if isinstance(t1, torch.nn.parameter.Parameter):
              t1 = t1.data
          if isinstance(t2, torch.nn.parameter.Parameter):
              t2 = t2.data
          self.assertEqual(type(t1), type(t2))
          if isinstance(t1, crypten.CrypTensor):
              t1 = t1.get_plain_text()
              t2 = t2.get_plain_text()
          self.assertTrue(t1.eq(t2).all())
  
      def _check_state_dict(model, state_dict):
          """
          Checks if state_dict matches parameters in model.
          """
          # get all parameters, buffers, and names from model:
          params_buffers = {}
          for func in ["named_parameters", "named_buffers"]:
              params_buffers.update(
                  {k: v for k, v in getattr(model, func)()}
              )
  
          # do all the checks:
          self.assertEqual(len(params_buffers), len(state_dict))
          for name, param_or_buffer in params_buffers.items():
              self.assertIn(name, state_dict)
              _check_equal(state_dict[name], param_or_buffer)
  
      # test for individual modules:
      module_args = {
          "BatchNorm1d": (400,),
          "BatchNorm2d": (3,),
          "BatchNorm3d": (6,),
          "Conv1d": (3, 6, 5),
          "Conv2d": (3, 6, 5),
          "Linear": (400, 120),
      }
      for module_name, args in module_args.items():
          for encrypt in [False, True]:
              # create module and get state dict:
              module = getattr(crypten.nn, module_name)(*args)
              if encrypt:
                  module.encrypt()
              state_dict = module.state_dict()
              _check_state_dict(module, state_dict)
  
              # load state dict into fresh module:
              new_module = getattr(crypten.nn, module_name)(*args)
              if encrypt:
                  with self.assertRaises(AssertionError):
                      new_module.load_state_dict(state_dict)
                  new_module.encrypt()
              new_module.load_state_dict(state_dict)
              _check_state_dict(new_module, state_dict)
  
              # check saving and loading from file for encrypted modules
              if encrypt:
                  f = io.BytesIO()
                  crypten.save(module.state_dict(), f)
                  f.seek(0)
                  new_module2 = getattr(crypten.nn, module_name)(*args)
                  new_module2.encrypt()
                  new_module2.load_state_dict(crypten.load(f))
                  _check_state_dict(new_module2, state_dict)
  
      # tests for model that is sequence of modules:
      for num_layers in range(1, 6):
          for encrypt in [False, True]:
              # some variables that we need:
              input_size = (3, 10)
              output_size = (input_size[0], input_size[1] - num_layers)
              layer_idx = range(input_size[1], output_size[1], -1)
  
              # construct sequential model:
              module_list = [
                  crypten.nn.Linear(num_feat, num_feat - 1)
                  for num_feat in layer_idx
              ]
              model = crypten.nn.Sequential(*module_list)
              if encrypt:
                  model.encrypt()
  
              # check state dict:
              state_dict = model.state_dict()
              _check_state_dict(model, state_dict)
  
              # load state dict into fresh model:
              state_dict = model.state_dict()
              module_list = [
                  crypten.nn.Linear(num_feat, num_feat - 1)
                  for num_feat in layer_idx
              ]
              new_model = crypten.nn.Sequential(*module_list)
              if encrypt:
                  with self.assertRaises(AssertionError):
                      new_model.load_state_dict(state_dict)
                  new_model.encrypt()
              new_model.load_state_dict(state_dict)
  
              # check new model:
              _check_state_dict(model, state_dict)
  
              # check saving and loading from file for encrypted modules
              if encrypt:
                  f = io.BytesIO()
                  crypten.save(model.state_dict(), f)
                  f.seek(0)
                  module_list = [
                      crypten.nn.Linear(num_feat, num_feat - 1)
                      for num_feat in layer_idx
                  ]
                  new_model2 = crypten.nn.Sequential(*module_list)
                  new_model2.encrypt()
                  new_model2.load_state_dict(crypten.load(f))
                  _check_state_dict(new_model2, state_dict)
  
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 23 for calling the TestNN obj's underlying member method code with example usage: obj.test_state_dict() without return types