Adding Documentation to Python Code

Documentation for Python is provided through docstrings and generated using Sphinx. Please reference the Google-style Python docstrings guide for docstring formatting examples.

Follow these instructions to document, generate, and publish a new Python docstring:

1. Add the docstring directly under the name of the target method. At a very minimum, please add descriptions of:

   • The method’s functional behavior

   • The arguments, as denoted by the Args section

   • The return value, as denoted by the Returns section

   • The exceptions that can be thrown (if applicable), as denoted by the Raises section

   Other sections such as Todo, Note, and Example should be added as needed.

   Here is an example Python docstring:

   def example_method(alignment: c_size_t, param: float) -> int:
       """
       This class is an example of how you can write docstrings.
       You can add multiple lines of those descriptions. Make sure to include
       useful information about your method.
   
       **Code Example:**
   
       .. code-block:: cpp
   
           // Here is a C++ code block
           std::vector<int32_t> foo(const std::vector<int32_t> &lst) {
               std::vector<int32_t> ret;
               for (const auto x : lst) {
                   ret.emplace_back(x * 2);
               }
               return ret;
           }
   
       And here is a verbatim-text diagram example:
   
       .. code-block:: text
   
           .------+---------------------------------.-----------------------------
           |            Block A (first)             |       Block B (second)
           +------+------+--------------------------+------+------+---------------
           | Next | Prev |   usable space           | Next | Prev | usable space..
           +------+------+--------------------------+------+--+---+---------------
           ^  |                                     ^         |
           |  '-------------------------------------'         |
           |                                                  |
           '----------- Block B's prev points to Block A -----'
   
       Todo:
           * This is a TODO item.
           * And a second TODO item.
   
       Args:
           alignment (c_size_t): Description of the `alignment` value.
           param (float): Description of `param1`.
   
       Returns:
           Description of the method's return value.
   
       Raises:
           AttributeError: If there is an error with the attributes.
           ValueError: If `param` is equal to 3.14.
   
       Example:
           This is how you can use this function
   
           >>> print("Code blocks are supported")
   
       Note:
           For more info on reStructuredText docstrings, see
           `here <https://www.sphinx-doc.org/en/master/usage/restructuredtext/basics.html>`__
           and
           `here <https://peps.python.org/pep-0287/>`__.
       """
       return 42
   
   
   

2. On the Sphinx documentation side, add an autofunction directive to the corresponding .rst file. If an .rst file for the corresponding Python source file does not exist, create a new one by the same name as the Python source file. Using the above example:

   .. autofunction:: fbgemm_gpu.docs.examples.example_method
   

3. Make sure the .rst file is included in to the toctree in index.rst (e.g. fbgemm-gpu.toc.api.python).

4. Verify the changes by building the docs locally with Building the Documentation or submitting a PR for a Netlify preview.

---

The Python docstring example above generates the following HTML output:

fbgemm_gpu.docs.examples.example_method(alignment: c_ulong, param: float) → int

This class is an example of how you can write docstrings. You can add multiple lines of those descriptions. Make sure to include useful information about your method.

Code Example:

// Here is a C++ code block
std::vector<int32_t> foo(const std::vector<int32_t> &lst) {
    std::vector<int32_t> ret;
    for (const auto x : lst) {
        ret.emplace_back(x * 2);
    }
    return ret;
}

And here is a verbatim-text diagram example:

.------+---------------------------------.-----------------------------
|            Block A (first)             |       Block B (second)
+------+------+--------------------------+------+------+---------------
| Next | Prev |   usable space           | Next | Prev | usable space..
+------+------+--------------------------+------+--+---+---------------
^  |                                     ^         |
|  '-------------------------------------'         |
|                                                  |
'----------- Block B's prev points to Block A -----'

Todo

• This is a TODO item.

• And a second TODO item.

Parameters:

• alignment (c_size_t) – Description of the alignment value.

• param (float) – Description of param1.

Returns:

Description of the method’s return value.

Raises:

• AttributeError – If there is an error with the attributes.

• ValueError – If param is equal to 3.14.

Example

This is how you can use this function

>>> print("Code blocks are supported")

Note

For more info on reStructuredText docstrings, see here and here.

---

Adding Documentation to Auto-Generated Python Code

Many FBGEMM_GPU Python API methods are auto-generated through PyTorch during the build process, and require docstrings to be attached after the fact. Follow these instructions to document auto-generated Python methods:

1. If needed, create a Python file under fbgemm_gpu/fbgemm_gpu/docs in the repo.

2. In the Python file, use the provided helper methods in fbgemm_gpu.docs.common to add attach a docstring to the target auto-generated method by method name. Here is an example from the codebase:

   import torch
   
   from .common import add_docs
   
   
   add_docs(
       torch.ops.fbgemm.jagged_2d_to_dense,
       """
   jagged_2d_to_dense(values, x_offsets, max_sequence_length) -> Tensor
   
   Converts a jagged tensor, with a 2D values array into a dense tensor, padding with zeros.
   
   Args:
       values (Tensor): 2D tensor containing the values of the jagged tensor.
   
       x_offsets (Tensor): 1D tensor containing the starting point of each jagged row in the values tensor.
   
       max_sequence_length (int): Maximum length of any row in the jagged dimension.
   
   Returns:
       Tensor: The padded dense tensor
   
   Example:
       >>> values = torch.tensor([[1,1],[2,2],[3,3],[4,4]])
       >>> x_offsets = torch.tensor([0, 1, 3])
       >>> torch.ops.fbgemm.jagged_2d_to_dense(values, x_offsets, 3)
       tensor([[[1, 1],
                [0, 0],
                [0, 0]],
               [[2, 2],
                [3, 3],
                [0, 0]]])
   
   """,
   )
   

3. If not already present, append the Python file to the imports list in fbgemm_gpu/fbgemm_gpu/docs/__init__.py. This will force the ad-hoc documentation to be loaded on fbgemm_gpu module load. For example:

   from . import the_new_doc_module
   

4. Follow the remaining steps in Adding Documentation to Python Code to render the docstring in the documentation.