Merge pull request #29520 from eashtianjr:eaj-patch-1

PiperOrigin-RevId: 252873111

tensorflow/python/training/warm_starting_util.py

@@ -50,7 +50,7 @@ class VocabInfo(
   See `tf.estimator.WarmStartSettings` for examples of using
   VocabInfo to warm-start.
 
-  Attributes:
+  Args:
     new_vocab: [Required] A path to the new vocabulary file (used with the model
       to be trained).
     new_vocab_size: [Required] An integer indicating how many entries of the new
@@ -70,8 +70,14 @@ class VocabInfo(
       linear weights for binary classification / regression).  An axis of 1
       could be used for warm-starting output layers with class vocabularies.
 
-      For example:
+  Returns:
+    A `VocabInfo` which represents the vocabulary information for warm-starting.
 
+  Raises:
+    ValueError: `axis` is neither 0 or 1.
+
+      Example Usage:
+```python
       embeddings_vocab_info = tf.VocabInfo(
           new_vocab='embeddings_vocab',
           new_vocab_size=100,
@@ -100,7 +106,8 @@ class VocabInfo(
           backup_initializer=tf.compat.v1.zeros_initializer(),
           axis=0)
 
-      Currently, only axis=0 and axis=1 are supported.
+      #Currently, only axis=0 and axis=1 are supported.
+  ```
   """
 
   def __new__(cls,