use lazy imports more consistently

PiperOrigin-RevId: 692116983

Author: tomvdw

tensorflow_datasets/core/download/download_manager.py

@@ -28,12 +28,11 @@
 from absl import logging
 from etils import epath
 from etils import epy
+from tensorflow_datasets.core.utils.lazy_imports_utils import promise
 from tensorflow_datasets.core.utils.lazy_imports_utils import tree
 
 with epy.lazy_imports():
   # pylint: disable=g-import-not-at-top
-  import promise
-
   from tensorflow_datasets.core import utils
   from tensorflow_datasets.core.download import checksums
   from tensorflow_datasets.core.download import downloader

tensorflow_datasets/core/download/extractor.py

@@ -29,11 +29,11 @@
 
 from absl import logging
 from etils import epath
-import promise
 from tensorflow_datasets.core import constants
 from tensorflow_datasets.core import utils
 from tensorflow_datasets.core.download import resource as resource_lib
 from tensorflow_datasets.core.utils import file_utils
+from tensorflow_datasets.core.utils.lazy_imports_utils import promise
 from tensorflow_datasets.core.utils.lazy_imports_utils import tensorflow as tf
 
 

tensorflow_datasets/robomimic/dataset_utils.py

@@ -22,8 +22,8 @@
 from etils import epath
 import h5py
 import numpy as np
+from tensorflow_datasets.core.utils.lazy_imports_utils import tree
 import tensorflow_datasets.public_api as tfds
-import tree
 
 
 

tensorflow_datasets/text/c4.py

@@ -25,9 +25,9 @@
 from absl import logging
 from etils import epath
 from tensorflow_datasets.core.utils.lazy_imports_utils import tensorflow as tf
+from tensorflow_datasets.core.utils.lazy_imports_utils import tree
 import tensorflow_datasets.public_api as tfds
 from tensorflow_datasets.text import c4_utils
-import tree
 
 PageFeatures = c4_utils.PageFeatures
 

tensorflow_datasets/translate/wmt.py

@@ -26,8 +26,8 @@
 
 from absl import logging
 from etils import epath
+from tensorflow_datasets.core.utils.lazy_imports_utils import tree
 import tensorflow_datasets.public_api as tfds
-import tree
 
 _DESCRIPTION = """\
 Translate dataset based on the data from statmt.org.