From 80e6cf4201ba557400e1bcc3fb7f0b1aaee0920e Mon Sep 17 00:00:00 2001
From: Tom van der Weide <weide@google.com>
Date: Fri, 1 Nov 2024 03:36:18 -0700
Subject: [PATCH] use lazy imports more consistently

PiperOrigin-RevId: 692116983
---
 tensorflow_datasets/core/download/download_manager.py | 3 +--
 tensorflow_datasets/core/download/extractor.py        | 2 +-
 tensorflow_datasets/robomimic/dataset_utils.py        | 2 +-
 tensorflow_datasets/text/c4.py                        | 2 +-
 tensorflow_datasets/translate/wmt.py                  | 2 +-
 5 files changed, 5 insertions(+), 6 deletions(-)

diff --git a/tensorflow_datasets/core/download/download_manager.py b/tensorflow_datasets/core/download/download_manager.py
index ac368fe4da6..21f1087c40b 100644
--- a/tensorflow_datasets/core/download/download_manager.py
+++ b/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
diff --git a/tensorflow_datasets/core/download/extractor.py b/tensorflow_datasets/core/download/extractor.py
index 32066bf6d3a..683b8511b32 100644
--- a/tensorflow_datasets/core/download/extractor.py
+++ b/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
 
 
diff --git a/tensorflow_datasets/robomimic/dataset_utils.py b/tensorflow_datasets/robomimic/dataset_utils.py
index 28c1a24b7f4..fdd7e495141 100644
--- a/tensorflow_datasets/robomimic/dataset_utils.py
+++ b/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
 
 
 
diff --git a/tensorflow_datasets/text/c4.py b/tensorflow_datasets/text/c4.py
index d45187883d8..9f2095bb7e4 100644
--- a/tensorflow_datasets/text/c4.py
+++ b/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
 
diff --git a/tensorflow_datasets/translate/wmt.py b/tensorflow_datasets/translate/wmt.py
index b6691fd8eca..d3077ddbd44 100644
--- a/tensorflow_datasets/translate/wmt.py
+++ b/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.