AI Skills
docstring

docstring

热门

按照 PyTorch 官方规范为 PyTorch 函数和方法编写 docstring（文档字符串）。在编写或更新 PyTorch 代码中的 docstring 时使用。

10万Star
2.9万Fork
更新于 2026/8/2
获取 Skill源代码

相关 Skills

gifgrep39万blucli38万recsys-pipeline-architect24万
SKILL.md
只读
名称
docstring
描述

按照 PyTorch 官方规范为 PyTorch 函数和方法编写 docstring（文档字符串）。在编写或更新 PyTorch 代码中的 docstring 时使用。

PyTorch Docstring 编写指南

本 Skill 介绍了如何遵循 torch/_tensor_docs.py 与 torch/nn/functional.py 中的约定规范，为 PyTorch 项目中的函数和方法编写 docstring。

通用原则

  • 所有 docstring 均使用 原始字符串（raw string，即 r"""..."""），以避免 LaTeX/数学公式中的反斜杠转义问题
  • 遵循 Sphinx/reStructuredText (reST) 文档格式
  • 做到 简洁而完整 —— 涵盖所有核心关键信息
  • 尽可能附带 示例代码（Examples）
  • 使用 交叉引用 关联相关函数/类

Docstring 结构

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 指令编写数学公式：

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

或行内公式：:math:x^2``

4. 交叉引用

使用 Sphinx 角色（roles）跳转关联的类或函数：

  • :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. 注意事项与警告

使用警告框（admonitions）标注重要提醒：

.. 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: 默认值”
  • 行内代码使用双反引号：``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 函数直接书写标准 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 语言绑定的底层函数（使用 _add_docstr）

对于底层 C/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)`
    ...
""",
)

就地操作版本（In-Place）

就地操作方法（以 _ 结尾）可直接引用原非就地版本：

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`
""",
)

常见写法模式

张量 Shape 标注

使用 LaTeX 数学公式语法标注张量的维度形状（Shape）：

: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 docstring 时，请确认：

  • [ ] 使用了原始字符串（r"""）
  • [ ] 首行包含了函数签名
  • [ ] 包含了简要功能描述
  • [ ] 在 Args 区域详细记录了所有参数及其类型
  • [ ] 为可选参数指明了默认值
  • [ ] 正确使用了 Sphinx 交叉引用（:func:, :class:, :meth:）
  • [ ] 如有需要，添加了数学公式
  • [ ] 在 Examples 区域至少提供了一个代码示例
  • [ ] 对关键注意事项补充了 warning/note 提醒
  • [ ] 使用 :class: 关联到了对应的模块类
  • [ ] 张量 Shape 标注规范采用了正确的数学公式表达
  • [ ] 整体缩进和格式保持一致

常用 Sphinx Roles 参考

  • :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 等