code: def test_save_restore_cycle(self, golden):
      module = golden.create_module()
  
      # Create all parameters and set them to sequential (but different) values.
      variables = golden.create_all_variables(module)
      for index, variable in enumerate(variables):
          variable.assign(goldens.range_like(variable, start=index))
  
      @tf.function(input_signature=[golden.input_spec])
      def inference(x):
          return golden.forward(module, x)
  
      # Create a saved model, add a method for inference and a dependency on our
      # module such that it can find dependencies.
      saved_model = snt.Module()
      saved_model._module = module
      saved_model.inference = inference
      saved_model.all_variables = list(module.variables)
  
      # Sample input.
      x = goldens.range_like(golden.input_spec)
  
      # Run the saved model and pull variable values.
      saved_model.inference(x)
      v1 = saved_model.all_variables
  
      # Save the model to disk and restore it.
      tmp_dir = os.path.join(
          absltest.get_default_test_tmpdir(), golden.name
      )
      tf.saved_model.save(saved_model, tmp_dir)
      restored_model = tf.saved_model.load(tmp_dir)
  
      # Run the loaded model and pull variable values.
      v2 = restored_model.all_variables
      y2 = restored_model.inference(x)
  
      if golden.deterministic:
          # The output from both the saved and restored model should be close.
          y1 = saved_model.inference(x)
          # TODO(b/161972382): The restored model doesn't seem to specialize the
          # graph with implementation selector, so the original model uses CuDNN
          # calls, whereas the restored model uses the non-specialized graph which
          # still contains a regular Tanh op.
          tree.map_structure(self.assertAllClose, y1, y2)
  
      for a, b in zip(v1, v2):
          self.assertEqual(a.name, b.name)
          self.assertEqual(a.device, b.device)
          self.assertAllEqual(a.read_value(), b.read_value())
  
file_path: sonnet/sonnet/src/conformance/saved_model_test.py
imports: import os
  from absl.testing import absltest
  from absl.testing import parameterized
  import sonnet
  from sonnet.src import test_utils
  from sonnet.src.conformance import goldens
  import tensorflow
  import tree
label_desc: Use image node_id 1 for calling the SavedModelTest obj's underlying member method code with example usage: obj.test_save_restore_cycle(golden) and returns: golden