Adding support for overriding do_not_doc_in_subclasses decorator. This is useful to document call methods on tf.keras.Model class and all its child classes.

Sequential won't document that method because of `do_not_doc_inheritable` decorator added in this CL.

PiperOrigin-RevId: 340943392
Change-Id: I80bcbce117f14eb098236d6cdf0fbeab0ad7e720

tensorflow/python/keras/engine/functional.py

@@ -46,6 +46,7 @@ from tensorflow.python.ops import math_ops
 from tensorflow.python.platform import tf_logging as logging
 from tensorflow.python.training.tracking import base as trackable
 from tensorflow.python.util import nest
+from tensorflow.tools.docs import doc_controls
 
 
 # pylint: disable=g-classes-have-attributes
@@ -403,6 +404,7 @@ class Functional(training_lib.Model):
     return nest.map_structure(lambda t: getattr(t, '_keras_mask', None),
                               output_tensors)
 
+  @doc_controls.do_not_doc_inheritable
   def call(self, inputs, training=None, mask=None):
     """Calls the model on new inputs.
 

tensorflow/python/keras/engine/training.py

@@ -422,6 +422,7 @@ class Model(base_layer.Layer, version_utils.ModelVersionSelector):
                            'the correct dtype).')
     super(Model, self).build(input_shape)
 
+  @doc_controls.doc_in_current_and_subclasses
   def call(self, inputs, training=None, mask=None):
     """Calls the model on new inputs.
 
@@ -429,6 +430,11 @@ class Model(base_layer.Layer, version_utils.ModelVersionSelector):
     all ops in the graph to the new inputs
     (e.g. build a new computational graph from the provided inputs).
 
+    Note: This method should not be called directly. It is only meant to be
+    overridden when subclassing `tf.keras.Model`.
+    To call a model on an input, always use the `__call__` method,
+    i.e. `model(inputs)`, which relies on the underlying `call` method.
+
     Arguments:
         inputs: A tensor or list of tensors.
         training: Boolean or boolean scalar tensor, indicating whether to run

tensorflow/tools/docs/doc_controls.py

@@ -286,3 +286,68 @@ def doc_private(obj: T) -> T:
 
   setattr(obj, _DOC_PRIVATE, None)
   return obj
+
+
+_DOC_IN_CURRENT_AND_SUBCLASSES = "_tf_docs_doc_in_current_and_subclasses"
+
+
+def doc_in_current_and_subclasses(obj: T) -> T:
+  """Overrides `do_not_doc_in_subclasses` decorator.
+
+  If this decorator is set on a child class's method whose parent's method
+  contains `do_not_doc_in_subclasses`, then that will be overriden and the
+  child method will get documented. All classes inherting from the child will
+  also document that method.
+
+  For example:
+
+  ```
+  class Parent:
+    @do_not_doc_in_subclasses
+    def method1(self):
+      pass
+    def method2(self):
+      pass
+
+  class Child1(Parent):
+    @doc_in_current_and_subclasses
+    def method1(self):
+      pass
+    def method2(self):
+      pass
+
+  class Child2(Parent):
+    def method1(self):
+      pass
+    def method2(self):
+      pass
+
+  class Child11(Child1):
+    pass
+  ```
+
+  This will produce the following docs:
+
+  ```
+  /Parent.md
+    # method1
+    # method2
+  /Child1.md
+    # method1
+    # method2
+  /Child2.md
+    # method2
+  /Child11.md
+    # method1
+    # method2
+  ```
+
+  Args:
+    obj: The class-attribute to hide from the generated docs.
+
+  Returns:
+    obj
+  """
+
+  setattr(obj, _DOC_IN_CURRENT_AND_SUBCLASSES, None)
+  return obj