code: def testLookupInNestedTFWhileLoop(self, source):
      strategy = parameter_server_strategy_v2.ParameterServerStrategyV2(
          self.cluster_resolver
      )
      coordinator = coordinator_lib.ClusterCoordinator(
          strategy=strategy
      )
  
      file_path = os.path.join(
          self.get_temp_dir(), "text_file_initializer"
      )
      with strategy.scope():
          model = self.Model(source, file_path)
  
      @def_function.function
      def replica_fn(batch_data):
          replica_result = array_ops.zeros(shape=(), dtype=dtypes.int64)
          for _ in math_ops.range(10):
              replica_result += math_ops.reduce_sum(
                  model.use_table(batch_data)
              )
          return replica_result
  
      @def_function.function
      def step_fn(iterator):
          step_result = array_ops.zeros(shape=(), dtype=dtypes.int64)
          for _ in math_ops.range(10):
              step_result += strategy.run(
                  replica_fn, args=(next(iterator),)
              )
  
          return step_result
  
      dataset = (
          dataset_ops.DatasetV2.from_tensors(
              constant_op.constant([0, 1, 3], dtype=dtypes.int64)
          )
          .repeat()
          .batch(24, drop_remainder=True)
          .prefetch(2)
      )
      distributed_dataset = coordinator.create_per_worker_dataset(
          strategy.experimental_distribute_dataset(dataset)
      )
  
      results = []
      for _ in range(10):
          results.append(
              coordinator.schedule(
                  step_fn, args=(iter(distributed_dataset),)
              )
          )
  
      coordinator.join()
  
      for r in results:
          self.assertAllClose(-2400, r.fetch())
  
imports: import copy
  import os
  from absl.testing import parameterized
  from tensorflow.python.compat import v2_compat
  from tensorflow.python.data.ops import dataset_ops
  from tensorflow.python.distribute import combinations
  from tensorflow.python.distribute import device_util
  from tensorflow.python.distribute import multi_process_runner
  from tensorflow.python.distribute import multi_worker_test_base
  from tensorflow.python.distribute import parameter_server_strategy_v2
  from tensorflow.python.distribute import ps_values
  from tensorflow.python.distribute.coordinator import cluster_coordinator
  from tensorflow.python.distribute.coordinator import coordinator_context
  from tensorflow.python.eager import def_function
  from tensorflow.python.eager import test
  from tensorflow.python.framework import constant_op
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import tensor_spec
  from tensorflow.python.module import module
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import lookup_ops
  from tensorflow.python.ops import math_ops
  from tensorflow.python.saved_model import load
  from tensorflow.python.saved_model import save
label_desc: Use image node_id 16 for calling the DistributedTableTest obj's underlying member method code with example usage: obj.testLookupInNestedTFWhileLoop(source) and returns: replica_result, step_result