code: def test_spmd_model_checkpointing(self):
      class LinearModel(module.Module):
          def __init__(self, w):
              super(LinearModel, self).__init__()
              self.w = variables.Variable(w)
  
          def __call__(self, x):
              return math_ops.matmul(x, self.w)
  
          def change_weights_op(self, w_new):
              return self.w.assign(w_new)
  
      batch_size = 32
      num_feature_in = 16
      num_feature_out = 8
      w1 = random_ops.random_uniform(
          (num_feature_in, num_feature_out), dtype=dtypes.float32
      )
      w2 = random_ops.random_uniform(
          (num_feature_in, num_feature_out), dtype=dtypes.float32
      )
      x = random_ops.random_uniform(
          (batch_size, num_feature_in), dtype=dtypes.float32
      )
  
      strategy, num_replicas = get_tpu_strategy(enable_spmd=True)
      with strategy.scope():
          model = LinearModel(w1)
  
      checkpoint_dir = self.get_temp_dir()
      checkpoint_prefix = os.path.join(checkpoint_dir, "ckpt")
      checkpoint = util.Checkpoint(model=model)
  
      @def_function.function
      def step_fn(x):
          x = strategy.experimental_split_to_logical_devices(x, [1, 2])
          return model(x)
  
      with self.cached_session() as sess:
          self.evaluate(variables.global_variables_initializer())
          checkpoint.save(file_prefix=checkpoint_prefix)
  
          self.evaluate(model.change_weights_op(w2))
          result = strategy.run(step_fn, args=(x,))
          self.assertAllClose(
              math_ops.matmul(x, w2) * num_replicas,
              self.evaluate(strategy.reduce("SUM", result, axis=None)),
              rtol=5e-3,
              atol=5e-3,
          )
  
          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(step_fn, args=(x,))
          self.assertAllClose(
              math_ops.matmul(x, w1) * num_replicas,
              self.evaluate(strategy.reduce("SUM", result, axis=None)),
              rtol=5e-3,
              atol=5e-3,
          )
  
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 10 for calling the TPUStrategyModelParallelismTest obj's underlying member method code with example usage: obj.test_spmd_model_checkpointing() and returns: model, math_ops, self