add comment explaining how to get global devices

keras-team · Sep 21, 2023 · f361c46 · f361c46
1 parent 2248912
commit f361c46
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/keras_core/backend/tensorflow/distribution_lib.py b/keras_core/backend/tensorflow/distribution_lib.py
@@ -25,6 +25,11 @@ def list_devices(device_type=None):
     device_type = (
         device_type.upper() if device_type else dtensor.preferred_device_type()
     )
+
+    # DTensor doesn't support getting global devices, even when knowing the
+    # Mesh. Use TF API instead to get global devices. Coordinator service is
+    # enabled by default with DTensor, so that list_logical_devices() returns
+    # a list of global devices. More context can be found in b/254911601.
     return tf.config.list_logical_devices(device_type=device_type)