Cherrypick DTensor docstring fix for 2.9 release. (#16434)

* Enable the keras dtensor API in OSS.

PiperOrigin-RevId: 438858608

* Switching learning/brain dependency to OSS compatible test_util

This is one test file failing, due to the monkey patching happens in the dtensor.init(), and I will need to dig more about the root cause (probably due to patching tf.Variable with DVariable, and cause logic difference for instance type checking.)

PiperOrigin-RevId: 439676157

* Update the docstring for keras.dtensor components.

1. Add docstring for LayoutMap.
2. Hide certain methods for keras.dtensor.optimizers.

PiperOrigin-RevId: 442651549

Author: qlzh727

keras/dtensor/layout_map.py

@@ -46,18 +46,41 @@ def get_current_layout_map():
 
 @keras_export('keras.dtensor.experimental.LayoutMap', v1=[])
 class LayoutMap(collections.abc.MutableMapping):
+  """A dict-like object that maps string to `Layout` instances.
 
-  def __init__(self, mesh=None):
-    """A dict like object that maps between string name and dtensor.Layout.
+  `LayoutMap` uses a string as key and a `Layout` as value. There is a behavior
+  difference between a normal Python dict and this class. The string key will be
+  treated as a regex when retrieving the value. See the docstring of
+  `get` for more details.
 
-    Note that this class might behave differently than a normal dict, eg, it
-    will treat the all the existing keys as a regex to map against input key.
+  See below for a usage example. You can define the naming schema
+  of the `Layout`, and then retrieve the corresponding `Layout` instance.
 
-    Args:
-      mesh: An optional dtensor.Mesh that is used to provide all replicated
-        layout as default when there isn't a layout is found based on the
-        mapping.
-    """
+  To use the `LayoutMap` with a `Model`, please see the docstring of
+  `tf.keras.dtensor.experimental.layout_map_scope`.
+
+  ```python
+  map = LayoutMap(mesh=None)
+  map['.*dense.*kernel'] = layout_2d
+  map['.*dense.*bias'] = layout_1d
+  map['.*conv2d.*kernel'] = layout_4d
+  map['.*conv2d.*bias'] = layout_1d
+
+  layout_1 = map['dense_1.kernel']    #   layout_1 == layout_2d
+  layout_2 = map['dense_1.bias']      #   layout_2 == layout_1d
+  layout_3 = map['dense_2.kernel']    #   layout_3 == layout_2d
+  layout_4 = map['dense_2.bias']      #   layout_4 == layout_1d
+  layout_5 = map['my_model/conv2d_123/kernel']    #   layout_5 == layout_4d
+  layout_6 = map['my_model/conv2d_123/bias']      #   layout_6 == layout_1d
+  ```
+
+  Args:
+    mesh: An optional `Mesh` that can be used to create all replicated
+      layout as default when there isn't a layout found based on the input
+      string query.
+  """
+
+  def __init__(self, mesh=None):
     self._layout_map = collections.OrderedDict()
     self._default_mesh = mesh
 
@@ -105,9 +128,17 @@ def __iter__(self):
     return iter(self._layout_map)
 
   def get_default_mesh(self):
+    """Return the default `Mesh` set at instance creation.
+
+    The `Mesh` can be used to create default replicated `Layout` when there
+    isn't a match of the input string query.
+    """
     return self._default_mesh
 
 
+LayoutMap.get.__doc__ = LayoutMap.__getitem__.__doc__
+
+
 @keras_export('keras.dtensor.experimental.layout_map_scope', v1=[])
 @contextlib.contextmanager
 def layout_map_scope(layout_map):

keras/dtensor/optimizers.py

@@ -26,6 +26,7 @@
 import tensorflow.compat.v2 as tf
 
 from tensorflow.python.util.tf_export import keras_export  # pylint: disable=g-direct-tensorflow-import
+from tensorflow.tools.docs import doc_controls
 
 
 # pylint: disable=protected-access,missing-class-docstring
@@ -103,6 +104,7 @@ def add_variable_from_reference(self,
         dtype=model_variable.dtype,
         trainable=False)
 
+  @doc_controls.do_not_generate_docs
   def aggregate_gradients(self, grads_and_vars):
     # Hide the aggregate_gradients from Optimizer.aggregate_gradients
     raise NotImplementedError(