update distribution_lib files docstrings. (keras-team#20625)

SamanehSaadat · Dec 10, 2024 · 1c70a85 · 1c70a85
1 parent 8968ea8
commit 1c70a85
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 12 deletions.
diff --git a/keras/src/backend/jax/distribution_lib.py b/keras/src/backend/jax/distribution_lib.py
@@ -1,10 +1,4 @@
-"""!!!DO NOT USE!!!
-
-Distribution related class for JAX backend.
-
-This is just a prototype and we might want to unify it
-with other backends in the future.
-"""
+"""Utilities for distribution strategy with JAX backend."""
 
 import jax
 import numpy as np

diff --git a/keras/src/distribution/distribution_lib.py b/keras/src/distribution/distribution_lib.py
@@ -1,9 +1,7 @@
-"""Unified high level distribution APIs across backends.
+"""Unified high-level distribution APIs across backends.
 
-!!!DO NOT USE!!! Currently under development and APIs are not final.
-
-Currently only the JAX backend has been implemented. The TensorFlow backend
-will be implemented in the future (via tf.dtensor API).
+Currently only the JAX backend is supported. The TensorFlow backend
+will be supported in the future (via tf.dtensor API).
 """
 
 import collections