code: def test_paritioned_model_checkpointing(self):
      class PartitionedModel(module.Module):
          def __init__(self, v, w):
              super(PartitionedModel, self).__init__()
  
              assert distribute_lib.has_strategy()
              strategy = distribute_lib.get_strategy()
  
              with strategy.extended.experimental_logical_device(0):
                  self.v = variables.Variable(v)
              with strategy.extended.experimental_logical_device(1):
                  self.w = variables.Variable(w)
  
          def __call__(self, x):
              replica_ctx = distribute_lib.get_replica_context()
              with replica_ctx.experimental_logical_device(0):
                  y = self.v * x
              with replica_ctx.experimental_logical_device(1):
                  z = self.w * y
              return z
  
          def change_weights_op(self, v_new, w_new):
              return control_flow_ops.group(
                  [self.v.assign(v_new), self.w.assign(w_new)]
              )
  
      strategy, num_replicas = get_tpu_strategy()
      with strategy.scope():
          model = PartitionedModel(2.0, 3.0)
  
      checkpoint_dir = self.get_temp_dir()
      checkpoint_prefix = os.path.join(checkpoint_dir, "ckpt")
      checkpoint = util.Checkpoint(model=model)
  
      with self.cached_session() as sess:
          self.evaluate(variables.global_variables_initializer())
          checkpoint.save(file_prefix=checkpoint_prefix)
  
          self.evaluate(model.change_weights_op(1.0, 4.0))
          result = strategy.run(
              def_function.function(model), args=(5.0,)
          )
          self.assertEqual(
              20.0 * num_replicas,
              self.evaluate(strategy.reduce("SUM", result, axis=None)),
          )
  
          status = checkpoint.restore(
              checkpoint_management.latest_checkpoint(checkpoint_dir)
          )
          status.run_restore_ops(
              sess
          )  # must run restore op in non-eager mode.
          status.assert_consumed()
          status.assert_existing_objects_matched()
          result = strategy.run(
              def_function.function(model), args=(5.0,)
          )
          self.assertEqual(
              30.0 * num_replicas,
              self.evaluate(strategy.reduce("SUM", result, axis=None)),
          )
  
imports: import os
  from absl.testing import parameterized
  from tensorflow.python.checkpoint import checkpoint
  from tensorflow.python.checkpoint import checkpoint_management
  from tensorflow.python.compiler.xla.experimental import xla_sharding
  from tensorflow.python.distribute import distribute_lib
  from tensorflow.python.distribute import packed_distributed_variable
  from tensorflow.python.distribute import strategy_test_lib
  from tensorflow.python.distribute import tpu_replicated_variable
  from tensorflow.python.distribute import tpu_strategy
  from tensorflow.python.distribute import tpu_values
  from tensorflow.python.distribute.cluster_resolver import tpu_cluster_resolver
  from tensorflow.python.eager import def_function
  from tensorflow.python.eager import remote
  from tensorflow.python.eager import test
  from tensorflow.python.framework import config
  from tensorflow.python.framework import constant_op
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import ops
  from tensorflow.python.framework import summary_test_util
  from tensorflow.python.module import module
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import control_flow_ops
  from tensorflow.python.ops import math_ops
  from tensorflow.python.ops import nn
  from tensorflow.python.ops import random_ops
  from tensorflow.python.ops import summary_ops_v2
  from tensorflow.python.ops import variable_scope
  from tensorflow.python.ops import variables
  from tensorflow.python.platform import flags
  from tensorflow.python.tpu import device_assignment
  from tensorflow.python.tpu import tpu_replication
label_desc: Use image node_id 3 for calling the TPUStrategyModelParallelismTest obj's underlying member method code with example usage: obj.test_paritioned_model_checkpointing() and returns: z, control_flow_ops