class_docstr: This class tests the crypten.nn package.
class_docstr_tok: ["This","class","tests","the","crypten.nn","package","."]
docstr: Test BatchNorm modules correctly set and update running stats. Also
  tests stateless BatchNormalization module.
code: def test_batchnorm_module(self) -> None:
      """
      Test BatchNorm modules correctly set and update running stats. Also
      tests stateless BatchNormalization module.
      """
      batchnorm_fn_and_size = (
          ("BatchNorm1d", (100, 5, 4)),
          ("BatchNorm2d", (100, 7, 4, 20)),
          ("BatchNorm3d", (100, 5, 4, 8, 15)),
      )
      for batchnorm_fn, size in batchnorm_fn_and_size:
          for is_training in (True, False):
              # create random input tensor:
              tensor = get_random_test_tensor(size=size, is_float=True)
              tensor.requires_grad = True
              encrypted_input = crypten.cryptensor(
                  tensor, requires_grad=True
              )
  
              # sample random weight and bias:
              C = size[1]
              params = {
                  "weight": get_random_test_tensor(
                      size=[C], max_value=1, is_float=True
                  ),
                  "bias": get_random_test_tensor(
                      size=[C], max_value=1, is_float=True
                  ),
              }
              for param in params.values():
                  param.requires_grad = True
  
              # dimensions for mean and variance:
              stats_dimensions = list(range(tensor.dim()))
              # perform on C dimension for tensor of shape (N, C, +)
              stats_dimensions.pop(1)
  
              # create models:
              enc_model = getattr(crypten.nn.module, batchnorm_fn)(C)
              plain_model = getattr(torch.nn.modules, batchnorm_fn)(C)
              for key, param in params.items():
                  enc_model.set_parameter(key, param)
                  setattr(plain_model, key, torch.nn.Parameter(param))
              enc_model = enc_model.encrypt()
  
              # check initial running statistics:
              stats = ["running_mean", "running_var"]
              for stat in stats:
                  self._check(
                      enc_model._buffers[stat],
                      plain_model._buffers[stat],
                      f"{stat} initial module value incorrect with train={is_training}",
                  )
              orig_buffers = copy.deepcopy(plain_model._buffers)
  
              # set training mode
              plain_model.train(is_training)
              enc_model.train(is_training)
  
              # check output and running_stats update:
              encr_output = enc_model.forward(encrypted_input)
              output = plain_model.forward(tensor)
              self._check(
                  encr_output,
                  output,
                  f"output of module incorrect with train={is_training}",
              )
              for stat in stats:
                  self._check(
                      enc_model._buffers[stat],
                      plain_model._buffers[stat],
                      f"{stat} momentum update in module incorrect with train={is_training}",
                  )
  
              # compute output of stateless batchnorm module:
              stateless_model = (
                  crypten.nn.module.BatchNormalization().encrypt()
              )
              stateless_model.train(is_training)
              encr_output = stateless_model(
                  (
                      encrypted_input,
                      crypten.cryptensor(params["weight"]),
                      crypten.cryptensor(params["bias"]),
                      crypten.cryptensor(orig_buffers["running_mean"]),
                      crypten.cryptensor(orig_buffers["running_var"]),
                  )
              )
              if not is_training:
                  encr_output = (encr_output,)
  
              # check that output of stateless batchnorm module is correct:
              self.assertEqual(
                  len(encr_output), 5 if is_training else 1
              )
              self._check(
                  encr_output[0],
                  output,
                  f"output of BatchNormalization incorrect with train={is_training}",
              )
              if (
                  is_training
              ):  # statistics are only updated at training time
                  for idx, stat in enumerate(stats):
                      self._check(
                          encr_output[1 + idx],
                          plain_model._buffers[stat],
                          f"{stat} update in BatchNormalization module incorrect",
                      )
  
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 21 for calling the TestNN obj's underlying member method code with example usage: obj.test_batchnorm_module() without return types