docstring

PyTorch Docstring 编写指南

本技能介绍如何为 PyTorch 项目中的函数和方法编写文档字符串，并遵循 torch/_tensor_docs.py 和 torch/nn/functional.py 中的约定。

通用原则

• 所有文档字符串均使用原始字符串（r"""..."""），以避免 LaTeX／数学公式反斜杠带来的问题
• 文档采用 Sphinx/reStructuredText（reST）格式
• 简洁但完整——包含所有必要信息
• 尽可能始终提供示例
• 使用交叉引用链接到相关函数／类

文档字符串结构

1. 函数签名（第一行）

以函数签名开头，展示所有参数：

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

注意：

• 包含函数名称
• 展示位置参数和仅限关键字参数（使用 * 分隔符）
• 包含默认值
• 展示返回类型注解
• 该行末尾不应有句号

2. 简要描述

用一行描述函数的作用：

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. 数学公式（如适用）

使用 Sphinx 数学指令表示数学表达式：

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

或使用行内数学公式：:math:\x^2``

4. 交叉引用

使用 Sphinx 角色链接到相关类和函数：

• :class:\~torch.nn.ModuleName`` - 链接到类
• :func:\torch.function_name`` - 链接到函数
• :meth:\~Tensor.method_name`` - 链接到方法
• :attr:\attribute_name`` - 引用属性
• 前缀 ~ 仅显示最后一个部分（例如，显示 Conv2d，而非 torch.nn.Conv2d）

示例：

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

5. 注释和警告

对重要信息使用提示块：

.. 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 部分

为所有参数编写文档，并注明类型和描述：

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

格式规则：

• 参数名称使用小写
• 类型写在括号中：(Type)；可选参数写作 (Type, optional)
• 类型后接描述
• 对于可选参数，在末尾注明“Default: value”
• 行内代码使用双反引号：``None``
• 续行缩进 2 个空格

7. Keyword Args 部分（如适用）

有时会单独记录关键字参数：

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 部分（如需要）

记录返回值：

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`.

或者，如果上下文足够明确，也可以直接在函数签名行中注明返回值。

9. Examples 部分

尽可能始终提供示例：

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)

格式规则：

• 使用双冒号 Examples::
• Python 代码使用 >>> 提示符
• 适当时使用 # 添加注释
• 有助于理解时展示实际输出（缩进但不带 >>>）

10. 外部引用

链接到论文或外部文档：

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

在文本中引用：See `Link Name`_

方法类型

原生 Python 函数

对于常规 Python 函数，使用标准文档字符串：

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 绑定函数（使用 add_docstr）

对于 C 绑定函数，使用 _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: input tensor of shape :math:`(\text{minibatch} , \text{in\_channels} , iW)`
    weight: filters of shape :math:`(\text{out\_channels} , kW)`
    ...
""",
)

原地变体

对于原地操作（名称以 _ 结尾），引用原始函数：

add_docstr_all(
    "abs_",
    r"""
abs_() -> Tensor

In-place version of :meth:`~Tensor.abs`
""",
)

别名函数

对于别名，直接引用原始函数：

add_docstr_all(
    "absolute",
    r"""
absolute() -> Tensor

Alias for :func:`abs`
""",
)

常见模式

形状说明

使用 LaTeX 数学记法表示张量形状：

:math:`(\text{minibatch} , \text{in\_channels} , iH , iW)`

可复用的参数定义

对于常用参数，定义一次并重复使用：

common_args = parse_kwargs(
    """
    dtype (:class:`torch.dtype`, optional): the desired type of returned tensor.
        Default: if None, same as this tensor.
"""
)

# Then use with .format():
r"""
...

Keyword args:
    {dtype}
    {device}
""".format(**common_args)

模板插入

插入可复现性说明或其他常用文本：

r"""
{tf32_note}

{cudnn_reproducibility_note}
""".format(**reproducibility_notes, **tf32_notes)

完整示例

以下是展示所有要素的完整示例：

def gumbel_softmax(
    logits: Tensor,
    tau: float = 1,
    hard: bool = False,
    eps: float = 1e-10,
    dim: int = -1,
) -> Tensor:
    r"""
    Sample from the Gumbel-Softmax distribution and optionally discretize.

    Args:
        logits (Tensor): `[..., num_features]` unnormalized log probabilities
        tau (float): non-negative scalar temperature
        hard (bool): if ``True``, the returned samples will be discretized as one-hot vectors,
              but will be differentiated as if it is the soft sample in autograd. Default: ``False``
        dim (int): A dimension along which softmax will be computed. Default: -1

    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`.

    .. note::
        This function is here for legacy reasons, may be removed from nn.Functional in the future.

    Examples::
        >>> logits = torch.randn(20, 32)
        >>> # Sample soft categorical using reparametrization trick:
        >>> F.gumbel_softmax(logits, tau=1, hard=False)
        >>> # Sample hard categorical using "Straight-through" trick:
        >>> F.gumbel_softmax(logits, tau=1, hard=True)

    .. _Link 1:
        https://arxiv.org/abs/1611.00712
    """
    # implementation

快速检查清单

编写 PyTorch 文档字符串时，请确保：

• 使用原始字符串（r"""）
• 在第一行包含函数签名
• 提供简短描述
• 在 Args 部分列出所有参数及其类型
• 为可选参数注明默认值
• 使用 Sphinx 交叉引用（:func:、:class:、:meth:）
• 适用时添加数学公式
• 在 Examples 部分至少包含一个示例
• 为重要注意事项添加警告／说明
• 使用 :class: 链接到相关模块类
• 使用恰当的数学符号表示张量形状
• 遵循一致的格式和缩进

常见 Sphinx 角色参考

• :class:\~torch.nn.Module`` - 类引用
• :func:\torch.function`` - 函数引用
• :meth:\~Tensor.method`` - 方法引用
• :attr:\attribute`` - 属性引用
• :math:\equation`` - 行内数学公式
• :ref:\label`` - 内部引用
• ``code`` - 行内代码（使用双反引号）

补充说明

• 缩进：代码使用 4 个空格，参数描述续行使用 2 个空格
• 行长：尽可能将每行控制在 100 个字符以内
• 句号：句子以句号结尾，但签名行除外
• 反引号：代码使用双反引号：``True`` ``None`` ``False``
• 类型：常见类型包括 Tensor、int、float、bool、str、tuple、list 等。