Adding doc_controls.decorate_all_class_attributes to tensorflow_docs.

`_hide_layer_and_module` can be removed.

PiperOrigin-RevId: 324648586
Change-Id: I30559a92f02d69772fc2719cb69d36470b1f9fd7

tensorflow/tools/docs/generate2.py

@@ -151,28 +151,6 @@ class TfExportAwareVisitor(doc_generator_visitor.DocGeneratorVisitor):
     return (canonical_score,) + scores
 
 
-def _hide_layer_and_module_methods():
-  """Hide methods and properties defined in the base classes of keras layers."""
-  # __dict__ only sees attributes defined in *this* class, not on parent classes
-  module_contents = list(tf.Module.__dict__.items())
-  layer_contents = list(tf.keras.layers.Layer.__dict__.items())
-
-  for name, obj in module_contents + layer_contents:
-    if name == "__init__":
-      continue
-
-    if isinstance(obj, property):
-      obj = obj.fget
-
-    if isinstance(obj, (staticmethod, classmethod)):
-      obj = obj.__func__
-
-    try:
-      doc_controls.do_not_doc_in_subclasses(obj)
-    except AttributeError:
-      pass
-
-
 def build_docs(output_dir, code_url_prefix, search_hints=True):
   """Build api docs for tensorflow v2.
 
@@ -189,7 +167,11 @@ def build_docs(output_dir, code_url_prefix, search_hints=True):
     if not name.startswith("_"):
       doc_controls.hide_from_search(obj)
 
-  _hide_layer_and_module_methods()
+  for cls in [tf.Module, tf.keras.layers.Layer]:
+    doc_controls.decorate_all_class_attributes(
+        decorator=doc_controls.do_not_doc_in_subclasses,
+        cls=cls,
+        skip=["__init__"])
 
   try:
     doc_controls.do_not_generate_docs(tf.__operators__)