Use epy.lazy_imports in datasets instead of try/except.

PiperOrigin-RevId: 625597756

Author: marcenacp

tensorflow_datasets/__init__.py

@@ -40,6 +40,7 @@
 
 _TIMESTAMP_IMPORT_STARTS = time.time()
 from absl import logging
+from etils import epy
 import tensorflow_datasets.core.logging as _tfds_logging
 from tensorflow_datasets.core.logging import call_metadata as _call_metadata
 
@@ -55,7 +56,7 @@
   # pytype: disable=import-error
   # For builds that don't include all dataset builders, we don't want to fail on
   # import errors of dataset builders.
-  try:
+  with epy.lazy_imports():
     from tensorflow_datasets import audio
     from tensorflow_datasets import graphs
     from tensorflow_datasets import image
@@ -79,8 +80,6 @@
     from tensorflow_datasets import video
     from tensorflow_datasets import vision_language
 
-  except ImportError:
-    pass
   # pytype: enable=import-error
 
   _import_time_ms_dataset_builders = int(