code: def test_save_restore(self, golden, replicator_fn, use_function):
      replicator = self.replicator_or_skip(replicator_fn, use_function)
  
      with replicator.scope():
          module = golden.create_module()
          variables = golden.create_all_variables(module)
  
      def forward():
          per_replica = replicator.run(lambda: golden.forward(module))
          return tree.map_structure(
              lambda args: tf.stack(replicator.unwrap(args), axis=0),
              per_replica,
          )
  
      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)
  
      # Assign sequential values to the weights.
      for index, variable in enumerate(variables):
          variable.assign(goldens.range_like(variable, start=index))
  
      # Create a checkpoint and save the weights.
      checkpoint = TestCheckpoint(module=module)
      checkpoint.save()
  
      # Compute a forward pass of the previously saved module.
      before_save_ys = forward()
  
      # Assign different values into the weights and do another forward pass. The
      # result should be different.
      for variable in variables:
          variable.assign(-tf.ones_like(variable))
  
      if golden.deterministic:
          y = forward()
          self.assertNotAllClose(y, before_save_ys)
  
      # Restore from the checkpoint and assert the module is in the same state.
      checkpoint.restore_latest(assert_consumed=True)
  
      for index, variable in enumerate(variables):
          # Parameters should be restored to their previous values.
          self.assertAllEqual(
              variable.read_value(),
              goldens.range_like(variable, start=index),
              msg=variable.name,
          )
  
      if golden.deterministic:
          tree.map_structure(
              self.assertAllEqual, forward(), before_save_ys
          )
  
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 2 for calling the ReplicatorCheckpointTest obj's underlying member method code with example usage: obj.test_save_restore(golden, replicator_fn, use_function) and returns: tree