code: def testWithMultipleDevicesMirrored(self):
      if self.primary_device == "CPU":
          self.skipTest("No devices to mirror across.")
      elif self.primary_device == "GPU":
          devices = tf.config.experimental.list_logical_devices("GPU")
      else:
          devices = tf.config.experimental.list_logical_devices("TPU")
  
      strategy = replicator.Replicator(
          [device.name for device in devices]
      )
      with strategy.scope():
          mean_metric = TestMetric()
          var_metric = TestMetric()
          layer = batch_norm.CrossReplicaBatchNorm(
              False, False, mean_metric, var_metric
          )
  
      scale = tf.constant(0.5, shape=(5,))
      offset = tf.constant(2.0, shape=(5,))
  
      def foo():
          inputs = tf.random.normal([2, 3, 3, 5])
          outputs = layer(inputs, True, False, scale, offset)
          return inputs, outputs
  
      inputs, outputs = strategy.run(foo)
      local_mean_metric = strategy.experimental_local_results(
          mean_metric.value
      )
      local_var_metric = strategy.experimental_local_results(
          var_metric.value
      )
      self.assertAllEqual(
          local_mean_metric[0].numpy(), local_mean_metric[1].numpy()
      )
      self.assertAllEqual(
          local_var_metric[0].numpy(), local_var_metric[1].numpy()
      )
      mean = local_mean_metric[0]
      var = local_var_metric[0]
  
      for inp, out in zip(
          strategy.experimental_local_results(inputs),
          strategy.experimental_local_results(outputs),
      ):
          expected_out = (inp - mean) * tf.math.rsqrt(
              var + 1e-5
          ) * scale + offset
          self.assertAllClose(out, expected_out)
  
file_path: sonnet/sonnet/src/distribute/distributed_batch_norm_test.py
imports: from absl import logging
  from absl.testing import parameterized
  from sonnet.src import test_utils
  from sonnet.src.distribute import distributed_batch_norm
  from sonnet.src.distribute import replicator
  import tensorflow
label_desc: Use image node_id 2 for calling the CrossReplicaBatchNormTest obj's underlying member method code with example usage: obj.testWithMultipleDevicesMirrored() and returns: inputs,  outputs