code: def testWithTpuStrategy(self):
      if self.primary_device != "TPU":
          self.skipTest("TPU strategy only runs on TPU's")
  
      strategy = replicator.TpuReplicator()
      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,))
  
      @tf.function
      def run():
          def compute():
              inputs = tf.ones([2, 3, 3, 5])
              outputs = layer(inputs, True, False, scale, offset)
              return inputs, outputs
  
          return strategy.run(compute)
  
      inputs, outputs = run()
  
      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 3 for calling the CrossReplicaBatchNormTest obj's underlying member method code with example usage: obj.testWithTpuStrategy() and returns: strategy, inputs,  outputs