code: def test_restore_on_create_in_replica_context(
      self, golden, replicator_fn, use_function
  ):
      replicator = self.replicator_or_skip(replicator_fn, use_function)
  
      # Save a checkpoint from a non-distributed model.
      module = golden.create_module()
      normal_variables = golden.create_all_variables(module)
      for index, variable in enumerate(normal_variables):
          variable.assign(
              goldens.range_like(variable, start=(index + 1))
          )
      checkpoint = TestCheckpoint(module=module)
      checkpoint.save()
      golden.forward(module)
  
      with replicator.scope():
          module = golden.create_module()
  
      def forward():
          return replicator.run(lambda: golden.forward(module))
  
      if use_function:
          forward = tf.function(forward)
          if self.primary_device == "TPU":
              # TODO(b/132329316) Remove when `xla.compile` allows tf.device(TPU).
              forward = with_soft_placement(forward)
  
      checkpoint = TestCheckpoint(module=module)
      status = checkpoint.restore_latest(assert_consumed=False)
      result = forward()
      status.assert_consumed()
  
      if golden.deterministic:
          result_iter = iter(
              replicator.experimental_local_results(result)
          )
          first_replica = next(result_iter)
          for next_replica in result_iter:
              tree.map_structure(
                  self.assertAllEqual, first_replica, next_replica
              )
  
      if not golden.has_side_effects:
          replicator_variables = module.variables
          for normal, distributed in zip(
              normal_variables, replicator_variables
          ):
              self.assertAllClose(
                  normal.read_value(),
                  distributed.read_value(),
                  msg=normal.name,
              )
  
imports: import os
  from absl import logging
  from absl.testing import absltest
  from absl.testing import parameterized
  from sonnet.src import test_utils
  from sonnet.src.conformance import goldens
  from sonnet.src.distribute import replicator
  from sonnet.src.distribute import replicator_test_utils
  import tensorflow
  import tree
label_desc: Use image node_id 6 for calling the ReplicatorCheckpointTest obj's underlying member method code with example usage: obj.test_restore_on_create_in_replica_context(golden, replicator_fn, use_function) and returns: replicator