code: def test_restore_from_non_distributed(
      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()
  
      # Create the same model (new params) in the replicator scope.
      with replicator.scope():
          module2 = golden.create_module()
          replicator_variables = golden.create_all_variables(module2)
  
      # Ensure the distributed params are != the values in the checkpoint.
      for normal, distributed in zip(
          normal_variables, replicator_variables
      ):
          distributed.assign(tf.zeros_like(distributed))
          self.assertNotAllClose(
              normal.read_value(), distributed.read_value()
          )
  
      # Restore the checkpoint and ensure the parameters are the same.
      checkpoint = TestCheckpoint(module=module2)
      checkpoint.restore_latest(assert_consumed=True)
  
      for normal, distributed in zip(
          normal_variables, replicator_variables
      ):
          self.assertAllEqual(
              normal.read_value(),
              distributed.read_value(),
              msg=normal.name,
          )
  
      if golden.deterministic:
  
          def run_forward(module):
              forward = 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)
              return forward()
  
          y_before = run_forward(module)
          y_after = run_forward(module2)
          tree.map_structure(self.assertAllEqual, y_before, y_after)
  
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 4 for calling the ReplicatorCheckpointTest obj's underlying member method code with example usage: obj.test_restore_from_non_distributed(golden, replicator_fn, use_function) and returns: forward