Cleanup the legacy conditional import for dtensor in keras

PiperOrigin-RevId: 551267421
keras-team · Jul 26, 2023 · ab566fd · ab566fd
1 parent bdfb8aa
commit ab566fd
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 9 deletions.
diff --git a/keras/dtensor/BUILD b/keras/dtensor/BUILD
@@ -26,6 +26,9 @@ package(
 py_library(
  name = "dtensor",
  srcs = ["__init__.py"],
+ deps = [
+ "//:expect_tensorflow_installed",
+ ],
 )
 
 tf_py_test(

diff --git a/keras/dtensor/__init__.py b/keras/dtensor/__init__.py
@@ -14,13 +14,5 @@
 # ==============================================================================
 """Keras' DTensor library."""
 
-_DTENSOR_API_ENABLED = True
 
-
-# Conditional import the dtensor API, since it is currently broken in OSS.
-if _DTENSOR_API_ENABLED:
- from tensorflow.compat.v2.experimental import dtensor as dtensor_api
-else:
- # Leave it with a placeholder, so that the import line from other python
- # file will not break.
- dtensor_api = None
+from tensorflow.compat.v2.experimental import dtensor as dtensor_api