code: def testAccessingTableInStepFunction(self, source):
      strategy = parameter_server_strategy_v2.ParameterServerStrategyV2(
          self.cluster_resolver
      )
      coordinator = coordinator_lib.ClusterCoordinator(
          strategy=strategy
      )
  
      with strategy.scope():
          lookup_table = self.createStaticHashTable(
              init_source=source, vals=[0, 1, 2], default_value=-2
          )
  
      dataset = (
          dataset_ops.DatasetV2.from_tensors(
              constant_op.constant([0, 1, 3], dtype=dtypes.int64)
          )
          .repeat()
          .batch(24, drop_remainder=True)
          .prefetch(2)
      )
      dataset = dataset.map(lookup_table.lookup)
  
      distributed_dataset = strategy.experimental_distribute_dataset(
          dataset
      )
      distributed_dataset = coordinator.create_per_worker_dataset(
          distributed_dataset
      )
  
      @def_function.function
      def worker_fn(iterator):
          def replica_fn(inputs):
              return math_ops.reduce_sum(lookup_table.lookup(inputs))
  
          all_results = strategy.run(replica_fn, args=(next(iterator),))
          return all_results
  
      steps_per_epoch = 10
      distributed_iterator = iter(distributed_dataset)
      result = []
      for _ in range(steps_per_epoch):
          result.append(
              coordinator.schedule(
                  worker_fn, args=(distributed_iterator,)
              )
          )
  
      coordinator.join()
  
      for r in result:
          returned_input = r.fetch()
          self.assertAllClose(-24, returned_input)
  
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 12 for calling the DistributedTableTest obj's underlying member method code with example usage: obj.testAccessingTableInStepFunction(source) and returns: all_results, math_ops