pytorch-style

name	pytorch-style
description	PyTorch 编码风格 Skill。蒸馏自 pytorch/pytorch 源码、核心贡献者 PR review、 PyTorch 官方 Contributing Guide、Soumith Chintala 和 Adam Paszke 的设计决策记录。 触发词：「PyTorch 风格」「像 PyTorch 一样写 Python」「pytorch style」「深度学习库风格」。 适用：深度学习框架开发、科学计算库、Python + C++ 混合开发、autograd 系统。

PyTorch · 编码 DNA

"PyTorch is designed to be Pythonic first. If it feels awkward in Python, it's wrong." — Soumith Chintala

角色定义

此 Skill 激活后，你写出的代码应该让 PyTorch 核心 contributor 在 PR review 时感觉 「这代码是从 PyTorch 里长出来的」，而不是「这是 TensorFlow 代码转译过来的」。

这意味着：Pythonic 优先、API 直觉、数学意图显式。

---

命名 DNA

5 条直觉规则：

1. snake_case 统治一切，类名用 PascalCase

   # ✅
   class LinearLayer(nn.Module):
       def forward(self, input: Tensor) -> Tensor: ...
   
   def batch_norm(input, weight, bias, ...): ...
   
   # ❌
   class linearLayer(nn.Module): ...
   def batchNorm(input, weight, bias, ...): ...
   

2. Tensor 参数统一叫 input，输出叫 output（或语义名）

   # ✅ — 和数学符号对齐
   def relu(input: Tensor, inplace: bool = False) -> Tensor: ...
   def conv2d(input, weight, bias=None, stride=1, ...): ...
   
   # ❌ — 自造名字破坏一致性
   def relu(x: Tensor) -> Tensor: ...
   def conv2d(tensor, kernel, ...): ...
   

3. 布尔参数名用 is_ 或语义形容词，不用动词

   # ✅
   inplace: bool = False
   requires_grad: bool = True
   pin_memory: bool = False
   
   # ❌
   do_inplace: bool = False
   enable_grad: bool = True
   

4. 私有方法 _ 前缀，C++ 绑定方法 _C. 前缀

   class Tensor:
       def _make_subclass(cls, ...): ...        # 内部 Python
       def sum(self): return torch._C._sum(self)  # C++ 绑定
   

5. 数学符号在 docstring 里用 LaTeX，变量名尽量贴近论文符号

   def scaled_dot_product_attention(
       query: Tensor,   # Q 矩阵
       key: Tensor,     # K 矩阵
       value: Tensor,   # V 矩阵
   ) -> Tensor:
       r"""
       Computes scaled dot product attention:
       :math:`\text{Attention}(Q, K, V) = \text{softmax}(\frac{QK^T}{\sqrt{d_k}})V`
       """
   

---

结构偏好

nn.Module 设计哲学：

# ✅ 标准 Module 结构
class MultiHeadAttention(nn.Module):
    def __init__(
        self,
        embed_dim: int,
        num_heads: int,
        dropout: float = 0.0,
        bias: bool = True,
        *,              # keyword-only 参数分隔
        batch_first: bool = False,
    ) -> None:
        super().__init__()
        # 1. 保存超参数
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        # 2. 定义子模块
        self.out_proj = NonDynamicallyQuantizableLinear(embed_dim, embed_dim, bias=bias)
        # 3. 初始化权重
        self._reset_parameters()

    def _reset_parameters(self) -> None:
        """初始化逻辑单独抽出，方便子类覆盖"""
        nn.init.xavier_uniform_(self.out_proj.weight)

    def forward(
        self,
        query: Tensor,
        key: Tensor,
        value: Tensor,
    ) -> tuple[Tensor, Tensor | None]:
        ...

functional vs Module 的选择：

• nn.functional.xxx — 无状态操作，没有学习参数
• nn.Module — 有学习参数、有状态、需要 train()/eval() 切换
• 两者都提供：Module 内部调用 functional（PyTorch 的标准模式）

class Dropout(nn.Module):
    def forward(self, input: Tensor) -> Tensor:
        return F.dropout(input, self.p, self.training, self.inplace)  # Module 调 functional

类型注解：

• 所有公开 API 必须有完整类型注解
• 用 Tensor（不是 torch.Tensor），已导入
• 可选参数用 X | None（Python 3.10+ 风格）或 Optional[X]
• 返回类型永远标注，即使是 None

---

注释哲学

docstring 是 API 的一部分，和代码同等重要：

def embedding(
    input: Tensor,
    weight: Tensor,
    padding_idx: int | None = None,
    ...
) -> Tensor:
    r"""A simple lookup table that stores embeddings of a fixed dictionary and size.

    This module is often used to store word embeddings and retrieve them using indices.
    The input to the module is a list of indices, and the output is the corresponding
    word embeddings.

    Args:
        input (Tensor): LongTensor of arbitrary shape containing the indices to extract
        weight (Tensor): Float tensor of shape `(num_embeddings, embedding_dim)`
        padding_idx (int, optional): If specified, entries at :attr:`padding_idx` do
            not contribute to the gradient

    Returns:
        Tensor: of shape `(*input.shape, embedding_dim)`

    Example::

        >>> # Example of embedding usage
        >>> embedding_matrix = torch.randn(10, 3)
        >>> input = torch.LongTensor([[1, 2, 4, 5], [4, 3, 2, 9]])
        >>> F.embedding(input, embedding_matrix)
    """

注释规范：

• r""" 前缀（raw string）保留 LaTeX 反斜杠
• Args/Returns/Raises/Example 四段固定结构
• 数学公式用 :math:...`` 内联，或 .. math:: 块
• # type: ignore 注释要说明原因

---

反模式（绝不这样写）

1. 用 *args/**kwargs 吞掉所有参数 — 公开 API 必须显式签名

   # ❌
   def conv2d(*args, **kwargs): ...
   # ✅
   def conv2d(input, weight, bias=None, stride=1, padding=0, ...): ...
   

2. eager 模式里写 graph-style 代码 — 不要假设用户在 torch.jit.script 里

3. 直接 .data 访问绕过 autograd — 除非明确知道在做什么，不然是 bug

   # ❌ 危险，切断梯度追踪
   x.data = new_value
   # ✅
   with torch.no_grad():
       x.copy_(new_value)
   

4. inplace 操作不加 _ 后缀 — PyTorch 约定：inplace 版本用 _ 结尾

   x.add_(1)    # ✅ inplace，有下划线
   y = x.add(1) # ✅ 非 inplace，无下划线
   x.add(1)     # ❌ 不 inplace 却叫 add，令人困惑
   

5. 不检查 device 一致性 — 多 GPU 环境是常态

   # ✅
   if query.device != key.device:
       raise ValueError(f"query and key must be on the same device, "
                        f"got {query.device} and {key.device}")
   

6. 魔法维度假设 — 不注释 Tensor shape

   # ❌
   attn = torch.bmm(q, k.transpose(-2, -1))
   # ✅
   # q: (batch, seq, head_dim), k: (batch, seq, head_dim)
   # attn: (batch, seq, seq)
   attn = torch.bmm(q, k.transpose(-2, -1))
   

---

校验测试

1. Pythonic 测试：把这个 API 给一个熟悉 NumPy 的研究员看，他 10 秒内能猜出怎么用吗？
2. Shape 注释：每个关键 Tensor 操作有没有注释输入/输出 shape？
3. inplace 一致性：所有 inplace 操作都以 _ 结尾了吗？

---

来源

• pytorch/pytorch 源码（torch/nn/、torch/functional.py）
• PyTorch Contributing Guide
• Soumith Chintala 的 PyTorch 设计哲学演讲（NeurIPS 2019）
• Adam Paszke 博士论文：Automatic Differentiation in PyTorch