code: def helper_test_syncbn(self, x, bn_fn, syncbn_fn):
      # run regular batch norm in train and eval mode
      bn = bn_fn()
      yt = bn(x, training=True)
      ye = bn(x, training=False)
      # run replicated sync batch norm in train and eval mode
      sync_bn = syncbn_fn()
      sync_bn_train = objax.Parallel(
          lambda x: sync_bn(x, training=True), vc=sync_bn.vars()
      )
      sync_bn_eval = objax.Parallel(
          lambda x: sync_bn(x, training=False), vc=sync_bn.vars()
      )
      with sync_bn.vars().replicate():
          yt_syncbn = sync_bn_train(x)
          ye_syncbn = sync_bn_eval(x)
      # replicated sync bn should have the same behavior as non-replicated regular bn
      self.assertTensorsAlmostEqual(yt, yt_syncbn)
      self.assertTensorsAlmostEqual(ye, ye_syncbn)
      self.assertTensorsAlmostEqual(yt, yt_syncbn)
      self.assertTensorsAlmostEqual(
          bn.running_mean.value, sync_bn.running_mean.value
      )
      self.assertTensorsAlmostEqual(
          bn.running_var.value, sync_bn.running_var.value
      )
      # run replicated non-sync batch norm - it should yield different result
      non_sync_bn = bn_fn()
      non_sync_bn_train = objax.Parallel(
          lambda x: non_sync_bn(x, training=True), vc=non_sync_bn.vars()
      )
      non_sync_bn_eval = objax.Parallel(
          lambda x: non_sync_bn(x, training=False),
          vc=non_sync_bn.vars(),
      )
      with non_sync_bn.vars().replicate():
          yt_non_syncbn = non_sync_bn_train(x)
          ye_non_syncbn = non_sync_bn_eval(x)
      self.assertTensorsNotEqual(yt, yt_non_syncbn)
      self.assertTensorsNotEqual(ye, ye_non_syncbn)
  
imports: import os
  import unittest
  import jax.numpy
  import numpy
  import objax
label_desc: Use image node_id 3 for calling the TestSyncBatchnorm obj's underlying member method code with example usage: obj.helper_test_syncbn(x, bn_fn, syncbn_fn) without return types