code: def get_tpu_strategy(enable_spmd=False):
      resolver = get_tpu_cluster_resolver()
      remote.connect_to_cluster(resolver)
      topology = tpu_cluster_resolver.initialize_tpu_system(resolver)
      num_replicas = resolver.get_tpu_system_metadata().num_cores // 2
      device_assignment = device_assignment_lib.DeviceAssignment.build(
          topology,
          num_replicas=num_replicas,
          computation_shape=[1, 1, 1, 2],
      )
      strategy = tpu_lib.TPUStrategyV2(
          resolver,
          experimental_device_assignment=device_assignment,
          experimental_spmd_xla_partitioning=enable_spmd,
      )
      return strategy, num_replicas
  
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 2 for calling a global function with example usage: get_tpu_strategy(enable_spmd) and returns: strategy,  num_replicas