Disable AutoGraph on method that uses mangled names.

PiperOrigin-RevId: 258883792

Author: Dan Moldovan

tensorflow_datasets/core/__init__.py

@@ -14,8 +14,13 @@
 # limitations under the License.
 
 """API to define datasets."""
+# Ensure TensorFlow is importable and its version is sufficiently recent. This
+# needs to happen before anything else, since the imports below will try to
+# import tensorflow, too.
+from tensorflow_datasets.core import tf_compat
+tf_compat.ensure_tf_install()
 
-from tensorflow_datasets.core.dataset_builder import BeamBasedBuilder
+from tensorflow_datasets.core.dataset_builder import BeamBasedBuilder  # pylint:disable=g-import-not-at-top
 from tensorflow_datasets.core.dataset_builder import BuilderConfig
 from tensorflow_datasets.core.dataset_builder import DatasetBuilder
 from tensorflow_datasets.core.dataset_builder import GeneratorBasedBuilder

tensorflow_datasets/core/features/top_level_feature.py

@@ -39,13 +39,21 @@ def __init__(self, *args, **kwargs):
     self.__is_top_level = False
     super(TopLevelFeature, self).__init__(*args, **kwargs)
 
+  # AutoGraph doesn't support mangled names (__is_top_level), so we explicitly
+  # disable it in methods that use them, to avoid the warning.
+  # TODO(mdan): Remove decorator once AutoGraph supports mangled names.
+  @tf.autograph.experimental.do_not_convert()
   def _set_top_level(self):
     """Indicates that the feature is top level.
 
     Internal function called by `DatasetInfo`.
     """
     self.__is_top_level = True
 
+  # AutoGraph doesn't support mangled names (__is_top_level), so we explicitly
+  # disable it in methods that use them, to avoid the warning.
+  # TODO(mdan): Remove decorator once AutoGraph supports mangled names.
+  @tf.autograph.experimental.do_not_convert()
   def decode_example(self, serialized_example, decoders=None):
     # pylint: disable=line-too-long
     """Decode the serialize examples.

tensorflow_datasets/core/tf_compat.py

@@ -109,6 +109,9 @@ def _patch_for_tf1_13(tf):
           dataset_ops.get_legacy_output_shapes)
       tf.compat.v2.data.Dataset.output_types = property(
           dataset_ops.get_legacy_output_types)
+  if not hasattr(tf.autograph.experimental, "do_not_convert"):
+    tf.autograph.experimental.do_not_convert = (
+        tf.contrib.autograph.do_not_convert)
 
 
 def is_dataset(ds):