code: def test_sync_batch_normalization_forward_backward(
      seed,
      axis,
      decay_rate,
      eps,
      batch_stat,
      output_stat,
      ctx,
      func_name,
      comm_nccl_opts,
      no_scale,
      no_bias,
      no_mean,
      no_variance,
  ):
      if comm_nccl_opts is None:
          pytest.skip(
              "Communicator test is disabled. You can turn it on by an option `--test-communicator`."
          )
  
      comm = comm_nccl_opts.comm
      device_id = int(comm_nccl_opts.device_id)
      n_devices = len(comm_nccl_opts.devices)
      ctx.device_id = comm_nccl_opts.device_id
  
      def ref_batch_normalization(
          x,
          beta,
          gamma,
          rmean,
          rvar,
          comm,
          axes,
          decay_rate,
          eps,
          batch_stat,
          output_stat,
      ):
          orig = x - device_id
          inputs = []
          for i in range(n_devices):
              inputs.append(orig + i)
          x = np.concatenate(inputs)
  
          vx = nn.Variable(x.shape, True)
          vx.d = x
          vbeta = nn.Variable(beta.shape, True)
          vbeta.d = beta
          vgamma = nn.Variable(gamma.shape, True)
          vgamma.d = gamma
          vrmean = nn.Variable(rmean.shape, True)
          vrmean.d = rmean
          vrvar = nn.Variable(rvar.shape, True)
          vrvar.d = rvar
          with nn.context_scope(ctx):
              out = F.batch_normalization(
                  vx,
                  vbeta,
                  vgamma,
                  vrmean,
                  vrvar,
                  batch_stat=batch_stat,
                  output_stat=output_stat,
                  axes=axes,
                  decay_rate=decay_rate,
                  eps=eps,
              )
          if output_stat:
              out[0].forward()
              rmean[...] = vrmean.d.copy()
              rvar[...] = vrvar.d.copy()
              return (
                  out[0].d[device_id * 2 : (device_id + 1) * 2],
                  out[1].d,
                  out[2].d,
              )
          out.forward()
          rmean[...] = vrmean.d.copy()
          rvar[...] = vrvar.d.copy()
          return out.d[device_id * 2 : (device_id + 1) * 2]
  
      def ref_batch_normalize_grad(
          x,
          beta,
          gamma,
          rmean,
          rvar,
          dy,
          comm,
          axes,
          decay_rate,
          eps,
          batch_stat,
          output_stat,
          **kw,
      ):
          x_backup = x
          orig = x - device_id
          inputs = []
          for i in range(n_devices):
              inputs.append(orig + i)
          x = np.concatenate(inputs)
  
          vx = nn.Variable(x.shape, True)
          vx.d = x
          vx.g = 0
          vbeta = nn.Variable(beta.shape, True)
          vbeta.d = beta
          vbeta.g = 0
          vgamma = nn.Variable(gamma.shape, True)
          vgamma.d = gamma
          vgamma.g = 0
          vrmean = nn.Variable(rmean.shape, True)
          vrmean.d = rmean
          vrvar = nn.Variable(rvar.shape, True)
          vrvar.d = rvar
          with nn.context_scope(ctx):
              out, out_mean, out_var = F.batch_normalization(
                  vx,
                  vbeta,
                  vgamma,
                  vrmean,
                  vrvar,
                  batch_stat=batch_stat,
                  output_stat=True,
                  axes=axes,
                  decay_rate=decay_rate,
                  eps=eps,
              )
          f = out.parent
          f.forward(
              [vx, vbeta, vgamma, vrmean, vrvar],
              [out, out_mean, out_var],
          )
          for i in range(n_devices):
              out.g[2 * i : 2 * (i + 1)] = dy
          out_mean.grad.zero()
          out_var.grad.zero()
          f.backward(
              [vx, vbeta, vgamma, vrmean, vrvar],
              [out, out_mean, out_var],
          )
          worker_slice = np.s_[device_id * 2 : (device_id + 1) * 2]
          dx = vx.g[worker_slice].flatten()
  
          # Compute gradients of beta and gamma
          x = x_backup
          vx = nn.Variable.from_numpy_array(x, need_grad=False)
          with nn.context_scope(ctx), nn.auto_forward():
              out = F.batch_normalization(
                  vx,
                  vbeta,
                  vgamma,
                  out_mean,
                  out_var,
                  batch_stat=False,
                  output_stat=False,
                  axes=axes,
                  decay_rate=decay_rate,
                  eps=eps,
              )
          vbeta.grad.zero()
          vgamma.grad.zero()
          out.backward(dy, clear_buffer=True)
          return np.concatenate(
              [dx, vbeta.g.flatten(), vgamma.g.flatten()]
          )
  
      def ref_batch_normalize_grad_with_output_stat(
          x,
          beta,
          gamma,
          rmean,
          rvar,
          dy,
          dmean,
          dvar,
          comm,
          axes,
          decay_rate,
          eps,
          batch_stat,
          output_stat,
          **kw,
      ):
          x_backup = x
          orig = x - device_id
          inputs = []
          for i in range(n_devices):
              inputs.append(orig + i)
          x = np.concatenate(inputs)
  
          vx = nn.Variable(x.shape, True)
          vx.d = x
          vx.g = 0
          vbeta = nn.Variable(beta.shape, True)
          vbeta.d = beta
          vbeta.g = 0
          vgamma = nn.Variable(gamma.shape, True)
          vgamma.d = gamma
          vgamma.g = 0
          vrmean = nn.Variable(rmean.shape, True)
          vrmean.d = rmean
          vrvar = nn.Variable(rvar.shape, True)
          vrvar.d = rvar
          with nn.context_scope(ctx):
              out = F.batch_normalization(
                  vx,
                  vbeta,
                  vgamma,
                  vrmean,
                  vrvar,
                  batch_stat=batch_stat,
                  output_stat=True,
                  axes=axes,
                  decay_rate=decay_rate,
                  eps=eps,
              )
          out_mean, out_var = out[1:]
          f = out[0].parent
          f.forward([vx, vbeta, vgamma, vrmean, vrvar], out)
          for i in range(n_devices):
              out[0].g[2 * i : 2 * (i + 1)] = dy
          out[1].g[...] = dmean
          out[2].g[...] = dvar
          f.backward([vx, vbeta, vgamma, vrmean, vrvar], out)
          worker_slice = np.s_[device_id * 2 : (device_id + 1) * 2]
          dx = vx.g[worker_slice].flatten()
  
          # Compute gradients of beta and gamma
          x = x_backup
          vx = nn.Variable.from_numpy_array(x, need_grad=False)
          with nn.context_scope(ctx), nn.auto_forward():
              out = F.batch_normalization(
                  vx,
                  vbeta,
                  vgamma,
                  out_mean,
                  out_var,
                  batch_stat=False,
                  output_stat=False,
                  axes=axes,
                  decay_rate=decay_rate,
                  eps=eps,
              )
          vbeta.grad.zero()
          vgamma.grad.zero()
          out.backward(dy, clear_buffer=True)
          return np.concatenate(
              [dx, vbeta.g.flatten(), vgamma.g.flatten()]
          )
  
      from nbla_test_utils import function_tester
  
      rng = np.random.RandomState(seed)
      inputs = list(create_inputs(rng, axis, device_id))
      inputs = mask_inputs(
          inputs, no_scale, no_bias, no_mean, no_variance
      )
  
      if output_stat:
          ref_grad = ref_batch_normalize_grad_with_output_stat
      else:
          ref_grad = ref_batch_normalize_grad
      axes = [axis]
  
      # TODO: enable test using function_tester when the backward is supported with batch_stat=True
      if batch_stat is True:
          insert_identity = [True, True, True, False, False]
          function_tester(
              rng,
              F.sync_batch_normalization,
              ref_batch_normalization,
              inputs,
              ref_grad=ref_grad,
              func_args=[comm],
              func_kwargs=dict(
                  axes=axes,
                  decay_rate=decay_rate,
                  eps=eps,
                  batch_stat=batch_stat,
                  output_stat=output_stat,
              ),
              backward=[True, True, True, False, False],
              ctx=ctx,
              func_name=func_name,
              dstep=1e-2,
              atol_b=1e-2,
              insert_identity=insert_identity,
          )
      else:
          if no_mean or no_variance:
              return
  
          # Forward test when batch_stat is False
          vinputs = []
          for i in inputs:
              vinputs.append(nn.Variable(i.shape, True))
              vinputs[-1].d = i
  
          vinputs = mask_vinputs(
              vinputs, no_scale, no_bias, no_mean, no_variance
          )
  
          for i in range(5):
              inputs[0] = rng.randn(*inputs[0].shape) + device_id
              vinputs[0].d[...] = inputs[0]
              ref_y = ref_batch_normalization(
                  *(
                      inputs
                      + [
                          comm,
                          axes,
                          decay_rate,
                          eps,
                          batch_stat,
                          output_stat,
                      ]
                  )
              )
              with nn.context_scope(ctx), nn.auto_forward():
                  y = F.sync_batch_normalization(
                      *(
                          vinputs
                          + [
                              comm,
                              "world",
                              axes,
                              decay_rate,
                              eps,
                              batch_stat,
                              output_stat,
                          ]
                      )
                  )
              assert_allclose(vinputs[0].d, inputs[0])
  
      # Check if running mean and var works.
      if no_mean and no_variance:
          return
  
      vinputs = []
      for i in inputs:
          vinputs.append(nn.Variable(i.shape, True))
          vinputs[-1].d = i
  
      vinputs = mask_vinputs(
          vinputs, no_scale, no_bias, no_mean, no_variance
      )
  
      for i in range(5):
          inputs[0] = rng.randn(*inputs[0].shape) + device_id
          vinputs[0].d[...] = inputs[0]
          ref_y = ref_batch_normalization(
              *(
                  inputs
                  + [
                      comm,
                      axes,
                      decay_rate,
                      eps,
                      batch_stat,
                      output_stat,
                  ]
              )
          )
          with nn.context_scope(ctx), nn.auto_forward():
              y = F.sync_batch_normalization(
                  *(
                      vinputs
                      + [
                          comm,
                          "world",
                          axes,
                          decay_rate,
                          eps,
                          batch_stat,
                          output_stat,
                      ]
                  )
              )
          if not no_mean:
              assert_allclose(vinputs[3].d, inputs[3])
  
          if not no_variance:
              assert_allclose(vinputs[4].d, inputs[4], atol=1e-3)
  
file_path: nnabla/python/test/communicator/test_sync_batch_normalization.py
imports: import pytest
  import numpy
  import nnabla
  import nnabla.functions
  from nbla_test_utils import list_context
  from nnabla.testing import assert_allclose
  from six.moves import range
label_desc: Use image node_id 4 for calling a global function with example usage: test_sync_batch_normalization_forward_backward(seed, axis, decay_rate, eps, batch_stat, output_stat, ctx, func_name, comm_nccl_opts, no_scale, no_bias, no_mean, no_variance) and returns: out, np, np, out,  out,  out