/docstring

name: docstring
description: Write docstrings for PyTorch functions and methods following PyTorch conventions. Use when writing or updating docstrings in PyTorch code.

PyTorch Docstring Writing Guide

This skill describes how to write docstrings for functions and methods in the PyTorch project, following the conventions in `torch/_tensor_docs.py` and `torch/nn/functional.py`.

General Principles

• Use **raw strings** (`r"""..."""`) for all docstrings to avoid issues with LaTeX/math backslashes
• Follow **Sphinx/reStructuredText** (reST) format for documentation
• Be **concise but complete** - include all essential information
• Always include **examples** when possible
• Use **cross-references** to related functions/classes

Docstring Structure

1. Function Signature (First Line)

Start with the function signature showing all parameters:

r"""function_name(param1, param2, *, kwarg1=default1, kwarg2=default2) -> ReturnType

**Notes:**

• Include the function name
• Show positional and keyword-only arguments (use `*` separator)
• Include default values
• Show return type annotation
• This line should NOT end with a period

2. Brief Description

Provide a one-line description of what the function does:

r"""conv2d(input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1) -> Tensor

Applies a 2D convolution over an input image composed of several input
planes.

3. Mathematical Formulas (if applicable)

Use Sphinx math directives for mathematical expressions:

.. math::
    \text{Softmax}(x_{i}) = \frac{\exp(x_i)}{\sum_j \exp(x_j)}

Or inline math: `:math:\`x^2\``

4. Cross-References

Link to related classes and functions using Sphinx roles:

• `:class:\`~torch.nn.ModuleName\`` - Link to a class
• `:func:\`torch.function_name\`` - Link to a function
• `:meth:\`~Tensor.method_name\`` - Link to a method
• `:attr:\`attribute_name\`` - Reference an attribute
• The `~` prefix shows only the last component (e.g., `Conv2d` instead of `torch.nn.Conv2d`)

**Example:**

See :class:`~torch.nn.Conv2d` for details and output shape.

5. Notes and Warnings

Use admonitions for important information:

.. note::
    This function doesn't work directly with NLLLoss,
    which expects the Log to be computed between the Softmax and itself.
    Use log_softmax instead (it's faster and has better numerical properties).

.. warning::
    :func:`new_tensor` always copies :attr:`data`. If you have a Tensor
    ``data`` and want to avoid a copy, use :func:`torch.Tensor.requires_grad_`
    or :func:`torch.Tensor.detach`.

6. Args Section

Document all parameters with type annotations and descriptions:

Args:
    input (Tensor): input tensor of shape :math:`(\text{minibatch} , \text{in\_channels} , iH , iW)`
    weight (Tensor): filters of shape :math:`(\text{out\_channels} , kH , kW)`
    bias (Tensor, optional): optional bias tensor of shape :math:`(\text{out\_channels})`. Default: ``None``
    stride (int or tuple): the stride of the convolving kernel. Can be a single number or a
      tuple `(sH, sW)`. Default: 1

**Formatting rules:**

• Parameter name in **lowercase**
• Type in parentheses: `(Type)`, `(Type, optional)` for optional parameters
• Description follows the type
• For optional parameters, include "Default: ``value``" at the end
• Use double backticks for inline code: ``` ``None`` ```
• Indent continuation lines by 2 spaces

7. Keyword Args Section (if applicable)

Sometimes keyword arguments are documented separately:

Keyword args:
    dtype (:class:`torch.dtype`, optional): the desired type of returned tensor.
        Default: if None, same :class:`torch.dtype` as this tensor.
    device (:class:`torch.device`, optional): the desired device of returned tensor.
        Default: if None, same :class:`torch.device` as this tensor.
    requires_grad (bool, optional): If autograd should record operations on the
        returned tensor. Default: ``False``.

8. Returns Section (if needed)

Document the return value:

Returns:
    Tensor: Sampled tensor of same shape as `logits` from the Gumbel-Softmax distribution.
        If ``hard=True``, the returned samples will be one-hot, otherwise they will
        be probability distributions that sum to 1 across `dim`.

Or simply include it in the function signature line if obvious from context.

9. Examples Section

Always include examples when possible:

Examples::

    >>> inputs = torch.randn(33, 16, 30)
    >>> filters = torch.randn(20, 16, 5)
    >>> F.conv1d(inputs, filters)

    >>> # With square kernels and equal stride
    >>> filters = torch.randn(8, 4, 3, 3)
    >>> inputs = torch.randn(1, 4, 5, 5)
    >>> F.conv2d(inputs, filters, padding=1)

**Formatting rules:**

• Use `Examples::` with double colon
• Use `>>>` prompt for Python code
• Include comments with `#` when helpful
• Show actual output when it helps understanding (indent without `>>>`)

10. External References

Link to papers or external documentation:

.. _Link Name:
    https://arxiv.org/abs/1611.00712

Reference them in text: ```See `Link Name`_```

Method Types

Native Python Functions

For regular Python functions, use a standard docstring:

def relu(input: Tensor, inplace: bool = False) -> Tensor:
    r"""relu(input, inplace=False) -> Tensor

    Applies the rectified linear unit function element-wise. See
    :class:`~torch.nn.ReLU` for more details.
    """
    # implementation

C-Bound Functions (using add_docstr)

For C-bound functions, use `_add_docstr`:

conv1d = _add_docstr(
    torch.conv1d,
    r"""
conv1d(input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1) -> Tensor

Applies a 1D convolution over an input signal composed of several input
planes.

See :class:`~torch.nn.Conv1d` for details and output shape.

Args:
    input: