code: def check_multi_node_bn(
      comm, use_gpu=False, backend="auto", dtype=numpy.float32
  ):
      """Tests correctness of MultiNodeBatchNormalization.
  
      This test verifies MultiNodeBatchNormalization by comparing
      the following four configurations.
      (1) Single worker, normal BatchNormalization
      (2) Multiple workers, normal BatchNormalization
      (3) Single worker, MultiNodeBatchNormalization
      (4) Multiple workers, MultiNodeBatchNormalization
  
      Single worker: only using the result of worker 0, which uses the whole
          batch.
      Multiple workers: Each worker uses the 1/n part of the whole batch,
          where n is the number of nodes, and gradient is aggregated.
  
      This test conducts the forward and backward computation once for the
      deterministic model parameters and an input batch, and checks the
      gradients of parameters.
  
      The purpose of MultiNodeBatchNormalization is to make the results of
      (4) to be exactly same as (1). Therefore, the essential part is to
      check that the results of (1) and (4) are the same. The results of (3)
      should also be also same as them. In contrast, the results of (2) is
      not necessarily always same as them, and we can expect that it is
      almost always different. Therefore, we also check that the results of
      (2) is different from them, to see that this test working correctly.
      """
  
      local_batchsize = 8
      global_batchsize = local_batchsize * comm.size
      ndim = 3
      numpy.random.seed(71)
      x = numpy.random.random((global_batchsize, ndim)).astype(
          numpy.float32
      )
      y = numpy.random.randint(
          0, 1, size=global_batchsize, dtype=numpy.int32
      )
      x_local = comm.mpi_comm.scatter(
          x.reshape(comm.size, local_batchsize, ndim)
      )
      y_local = comm.mpi_comm.scatter(
          y.reshape(comm.size, local_batchsize)
      )
  
      io_dtype = dtype
      l_dtype = dtype
      bn_dtype = dtype
      if dtype == chainer.mixed16:
          io_dtype = numpy.float16
          l_dtype = numpy.float16
          bn_dtype = numpy.float32
  
      x = x.astype(io_dtype)
      x_local = x_local.astype(io_dtype)
  
      if use_gpu:
          x = chainer.cuda.to_gpu(x)
          y = chainer.cuda.to_gpu(y)
          x_local = chainer.cuda.to_gpu(x_local)
          y_local = chainer.cuda.to_gpu(y_local)
  
      cls = chainer.links.Classifier
      with chainer.using_config("dtype", dtype):
          # Single worker
          m1 = cls(ModelNormalBN())
          # Multi worker, Ghost BN
          m2 = cls(ModelNormalBN())
          # Single worker, MNBN
          m3 = cls(ModelDistributedBN(comm, backend=backend))
          # Multi worker, MNBN
          m4 = cls(ModelDistributedBN(comm, backend=backend))
          # NOTE: m1, m3 and m4 should behave in the same way.
          # m2 may be different.
  
      if use_gpu:
          device = cupy.cuda.Device()
          m1.to_device(device)
          m2.to_device(device)
          m3.to_device(device)
          m4.to_device(device)
  
      m2.copyparams(m1)
      m3.copyparams(m1)
      m4.copyparams(m1)
  
      l1 = m1(x, y)
      m1.cleargrads()
      l1.backward()
  
      l2 = m2(x_local, y_local)
      m2.cleargrads()
      l2.backward()
      comm.multi_node_mean_grad(m2)
  
      l3 = m3(x, y)
      m3.cleargrads()
      l3.backward()
  
      l4 = m4(x_local, y_local)
      m4.cleargrads()
      l4.backward()
      comm.multi_node_mean_grad(m4)
  
      if comm.rank == 0:
          for p1, p2, p3, p4 in zip(
              sorted(m1.namedparams()),
              sorted(m2.namedparams()),
              sorted(m3.namedparams()),
              sorted(m4.namedparams()),
          ):
              name = p1[0]
              assert p2[0] == name
              assert p3[0] == name
              assert p4[0] == name
  
              if "/l" in name:
                  param_dtype = l_dtype
              else:
                  param_dtype = bn_dtype
              assert p1[1].data.dtype == param_dtype
              assert p2[1].data.dtype == param_dtype
              assert p3[1].data.dtype == param_dtype
              assert p4[1].data.dtype == param_dtype
  
              assert_option = {"atol": 1e-4, "rtol": 1e-3}
              if param_dtype == numpy.float16:
                  assert_option = {"atol": 1e-2, "rtol": 1e-2}
  
              chainer.testing.assert_allclose(
                  p1[1].grad, p3[1].grad, **assert_option
              )
              chainer.testing.assert_allclose(
                  p1[1].grad, p4[1].grad, **assert_option
              )
  
              # This is to see that this test is valid.
              if comm.size >= 2:
                  assert_not_allclose(p1[1].grad, p2[1].grad)
  
imports: import chainer
  from chainer.backends.cuda import cupy
  import chainer.testing
  import chainer.utils
  import mpi4py.MPI
  import numpy
  import pytest
  from chainermn.communicators.naive_communicator import NaiveCommunicator
  from chainermn.communicators.pure_nccl_communicator import PureNcclCommunicator
  from chainermn.links import MultiNodeBatchNormalization
  from chainermn import nccl
label_desc: Use image node_id 1 for calling a global function with example usage: check_multi_node_bn(comm, use_gpu, backend, dtype) without return types