How to Write PyTorch Function Documentation with the docstring Skill
When contributing to PyTorch, writing docstrings that follow project conventions can be challenging. The rules span Sphinx/reStructuredText formatting, mathematical notation, type annotations, and PyTorch-specific patterns.
This article explains how to write convention-compliant docstrings using the docstring skill from the PyTorch repository.
What This Skill Does
The docstring skill guides creation of PyTorch-convention docstrings:
- Raw string (
r""") usage requirements - Function signature, argument, and return value formatting
- Sphinx cross-references (
:func:,:class:,:meth:) - LaTeX math notation (
.. math::,:math:) - Examples:: section authoring
Designed for developers contributing to PyTorch core or maintaining internal function documentation.
Installation
Prerequisites
- Claude Code installed
- Access to PyTorch repository source code
Install Command
claude mcp add github.com/pytorch/pytorch/tree/main/.claude/skills/docstring
Usage
Basic Usage
Specify the function that needs a docstring:
Write a PyTorch-convention docstring for this function
Docstring Structure
PyTorch docstrings follow this order:
- Function signature - with parameters and return type
- Brief description - one-line summary
- Math formulas (if applicable) -
.. math::directive - Args: - type and description for each parameter
- Returns: - return value description
- Examples:: - code examples with
>>>prompt
Example
r"""relu(input, inplace=False) -> Tensor
Applies the rectified linear unit function element-wise.
See :class:`~torch.nn.ReLU` for more details.
Args:
input (Tensor): input tensor
inplace (bool): operate in-place. Default: ``False``
Examples::
>>> input = torch.randn(2)
>>> torch.relu(input)
"""
Important Considerations
Use Raw Strings
Always use r""" for docstrings so LaTeX backslashes are processed correctly.
Double Backticks for Code
Use double backticks for inline code: ``True``, ``None``. Single backticks are reserved for Sphinx roles.
C-Bound Functions Use _add_docstr
For C++-implemented functions, docstrings are added via _add_docstr in Python code.
Summary
The docstring skill enables accurate creation of PyTorch-convention docstrings including Sphinx notation and LaTeX math. Use it when contributing documentation to PyTorch.
For full specifications, visit the skill detail page.