message: Remove the hack that we create an actual iterator to get its type_spec. Since we trace the dataset_fn, we can get its output's element_spec to get an iterator_spec.
diff: @@ -32,6 +32,8 @@ import threading
  import weakref
  from absl import logging
  from six.moves import queue
  +
  +from tensorflow.python.data.ops import iterator_ops
  from tensorflow.python.distribute import input_lib
  from tensorflow.python.distribute import parameter_server_strategy_v2
  from tensorflow.python.distribute.client import metric_utils
  @@ -1145,15 +1147,12 @@ class _PerWorkerDistributedDataset(object):
  per_worker_iterator = self._client._create_per_worker_resources(
  _create_per_worker_iterator)
  - # Create an iterator, so the consumer function of this iterator can start
  - # tracing using this iterator without needing to wait for the completion of
  - # the iterater creation. Note: the iterator shouldn't use memory until it is
  - # consumed.
  - # TODO(b/154675763): get rid of this workaround once we can make input_fn a
  - # tf.function.
  - iterator = _create_per_worker_iterator()
  + # Setting type_spec of each RemoteValue so that functions taking these
  + # RemoteValues as inputs can be traced.
  for iterator_remote_value in per_worker_iterator._values:
  - iterator_remote_value._set_type_spec(iterator._type_spec)
  + iterator_remote_value._set_type_spec(
  + iterator_ops.IteratorSpec(
  + self._dataset_fn.structured_outputs.element_spec))
  return _PerWorkerDistributedIterator(per_worker_iterator._values)
  @property
  