code: def test_spmd_with_summary(self):
      original_device_placement = config.get_soft_device_placement()
      config.set_soft_device_placement(True)
  
      strategy, _ = get_tpu_strategy(enable_spmd=True)
      summary_dir = self.get_temp_dir()
      writer = summary_ops.create_file_writer_v2(summary_dir)
      const_multiple = 2
      num_iters = 10
      expected_event_count = num_iters + 1
  
      with strategy.scope():
          step = variables.Variable(1, dtype=dtypes.int64)
  
      @def_function.function
      def run():
          with writer.as_default():
              with summary_ops.record_if(True):
                  summary_ops.scalar(
                      "result", step * const_multiple, step=step
                  )
                  step.assign_add(1)
  
      for _ in range(num_iters):
          strategy.run(run, args=())
  
      for val in step.values:
          for var in val.variables:
              self.assertAllEqual(expected_event_count, var)
  
      events = summary_test_util.events_from_logdir(summary_dir)
      self.assertLen(events, expected_event_count)
  
      # Event[0] is generic metadata and summary_ops data starts at event[1].
      for logged_step in range(1, expected_event_count):
          self.assertEqual(
              events[logged_step].summary.value[0].simple_value,
              logged_step * const_multiple,
          )
  
      config.set_soft_device_placement(original_device_placement)
  
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 11 for calling the TPUStrategyModelParallelismTest obj's underlying member method code with example usage: obj.test_spmd_with_summary() without return types