Change keras to use dtensor public API.

PiperOrigin-RevId: 438605222
keras-team · Mar 31, 2022 · 55476a8 · 55476a8
1 parent c9e2b4d
commit 55476a8
Showing 1 changed file with 1 addition and 4 deletions.
diff --git a/keras/dtensor/__init__.py b/keras/dtensor/__init__.py
@@ -19,11 +19,8 @@
 
 # Conditional import the dtensor API, since it is currently broken in OSS.
 if _DTENSOR_API_ENABLED:
-  # pylint: disable=g-direct-tensorflow-import, g-import-not-at-top
-  from tensorflow.dtensor import python as dtensor_api
+  from tensorflow.compat.v2.experimental import dtensor as dtensor_api  # pylint: disable=g-import-not-at-top
 else:
   # Leave it with a placeholder, so that the import line from other python file
   # will not break.
   dtensor_api = None
-
-