code: def testTypeSpecReturnedFromTFFunction(
      self, distribution, enable_get_next_as_optional
  ):
      # TODO(ishark): This is observed when tensor is copied from one device to
      # other and since DatasetVariantWrapper does not have a copy
      # function. Some Context: b/146981184
      # Try to renable with non-canonicalized input workers, which
      # helped in PS Strategy for similar error.
      self.skipTest(
          "Failures observed in Ubuntu presubmit: No unary variant  "
          "device copy function found for direction: 1 and Variant "
          "type_index:tensorflow::data::(anonymous namespace)::"
          "DatasetVariantWrapper"
      )
  
      @def_function.function
      def create_dist_dataset():
          dataset = dataset_ops.DatasetV2.range(10).batch(2)
          return distribution.experimental_distribute_dataset(dataset)
  
      distribution.extended.experimental_enable_get_next_as_optional = (
          enable_get_next_as_optional
      )
  
      dist_dataset = create_dist_dataset()
  
      spec = dist_dataset._type_spec
      self.assertEqual(spec._input_workers, dist_dataset._input_workers)
      self.assertEqual(
          spec._element_spec._value_specs,
          (
              tensor_spec.TensorSpec(
                  shape=(None,), dtype=dtypes.int64, name=None
              ),
              tensor_spec.TensorSpec(
                  shape=(None,), dtype=dtypes.int64, name=None
              ),
          ),
      )
  
      # Read distributed data to confirm values are correct.
      iterator = iter(dist_dataset)
      data = []
      for it in iterator:
          data.append(distribution.experimental_local_results(it))
      self.assertAllEqual(
          nest.flatten(data),
          list(
              dataset_ops.DatasetV2.range(10)
              .batch(1)
              .as_numpy_iterator()
          ),
      )
  
imports: import os
  from absl.testing import parameterized
  import numpy
  from tensorflow.python import tf2
  from tensorflow.python.data.ops import dataset_ops
  from tensorflow.python.data.ops import readers
  from tensorflow.python.distribute import combinations
  from tensorflow.python.distribute import distribute_lib
  from tensorflow.python.distribute import strategy_combinations
  from tensorflow.python.distribute import test_util
  from tensorflow.python.distribute import tpu_strategy
  from tensorflow.python.distribute import values
  from tensorflow.python.eager import def_function
  from tensorflow.python.eager import test
  from tensorflow.python.framework import composite_tensor
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import sparse_tensor
  from tensorflow.python.framework import tensor_shape
  from tensorflow.python.framework import tensor_spec
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import string_ops
  from tensorflow.python.ops.ragged import ragged_tensor
  from tensorflow.python.util import nest
label_desc: Use image node_id 2 for calling the DistributedDatasetTypeSpecTest obj's underlying member method code with example usage: obj.testTypeSpecReturnedFromTFFunction(distribution, enable_get_next_as_optional) and returns: distribution