code: def test_variable_creation_in_replica_context(
      self, golden, replicator_fn
  ):
      tf.random.set_seed(None)
      replicator = replicator_fn()
  
      with replicator.scope():
          mod = golden.create_module()
  
      @tf.function
      def forward():
          step = lambda: golden.create_all_variables(mod)
          return replicator.run(step)
  
      # TODO(b/132329316) Remove when `xla.compile` allows tf.device(TPU).
      with tf.device(None):
          variables_per_replica = forward()
  
      self.assertLen(variables_per_replica, golden.num_variables)
  
      for per_replica_variable in variables_per_replica:
          self.assertSameValuePerReplica(
              replicator, per_replica_variable
          )
  
file_path: sonnet/sonnet/src/conformance/distribute_test.py
imports: from typing import Callable, Tuple
  from absl.testing import parameterized
  import sonnet
  from sonnet.src import test_utils
  from sonnet.src.conformance import descriptors
  from sonnet.src.conformance import goldens
  from sonnet.src.distribute import replicator
  from sonnet.src.distribute import replicator_test_utils
  import tensorflow
label_desc: Use image node_id 1 for calling the TpuReplicatorTest obj's underlying member method code with example usage: obj.test_variable_creation_in_replica_context(golden, replicator_fn) and returns: replicator