请改用 jaxtyping

欢迎！对于新项目，我现在强烈建议使用我更新的 jaxtyping 项目。它支持 PyTorch，实际上并不依赖 JAX，而且与 TorchTyping 不同的是，它兼容静态类型检查器。名称中的“jax”现在只是历史遗留了！

原始的 torchtyping README 如下。

---

torchtyping

为张量的形状、数据类型、命名维度等添加类型注解

将这段代码：

def batch_outer_product(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    # x 的形状是 (batch, x_channels)
    # y 的形状是 (batch, y_channels)
    # 返回值的形状是 (batch, x_channels, y_channels)

    return x.unsqueeze(-1) * y.unsqueeze(-2)

改为：

def batch_outer_product(x:   TensorType["batch", "x_channels"],
                        y:   TensorType["batch", "y_channels"]
                        ) -> TensorType["batch", "x_channels", "y_channels"]:

    return x.unsqueeze(-1) * y.unsqueeze(-2)

同时对形状（数据类型等）规范进行程序化检查。

告别 bug！迎接强制执行且清晰明了的代码文档。

如果你（像我一样）经常在代码中写类似 # x 的形状是 (batch, hidden_state) 的注释，或者插入诸如 assert x.shape == y.shape 的语句，仅仅是为了跟踪各个张量的形状，那么这个库就非常适合你。

---

安装

pip install torchtyping

需要 Python >=3.7 和 PyTorch >=1.7.0。

如果使用 typeguard，则其版本必须低于 3.0.0。

使用方法

torchtyping 允许对以下内容进行类型注解：

• 形状：大小、维度数；
• 数据类型（浮点数、整数等）；
• 布局（稠密、稀疏）；
• 命名维度，基于 命名张量；
• 任意数量的批次维度，通过 ... 表示；
• 以及任何其他你想要的内容，因为 torchtyping 具有高度可扩展性。

如果（可选地）安装了 typeguard，则可以在运行时检查类型，以确保张量确实符合声明的形状、数据类型等。

# 示例

from torch import rand
from torchtyping import TensorType, patch_typeguard
from typeguard import typechecked

patch_typeguard()  # 在 @typechecked 之前调用

@typechecked
def func(x: TensorType["batch"],
         y: TensorType["batch"]) -> TensorType["batch"]:
    return x + y

func(rand(3), rand(3))  # 正常运行
func(rand(3), rand(1))
# TypeError: 维度 'batch' 的大小不一致。分别得到 1 和 3。

typeguard 还提供了一个导入钩子，可以用来自动测试整个模块，而无需手动添加 @typeguard.typechecked 装饰器。

如果你不使用 typeguard，则可以完全省略 torchtyping.patch_typeguard()，仅将 torchtyping 用于文档目的。如果你目前还没有在常规 Python 编程中使用 typeguard，强烈建议考虑使用它。这是一种很好的方式来消除 bug。此外，typeguard 和 torchtyping 都可以与 pytest 集成，因此如果你担心性能开销，也可以只在测试时启用它们。

API

torchtyping.TensorType[shape, dtype, layout, details]

该库的核心。

shape、dtype、layout、details 各参数均可选。

• shape 参数可以是以下任意一种：
  • 一个 int：维度必须精确为该大小。若为 -1，则允许任意大小。
  • 一个 str：运行时传入的维度大小将绑定到该名称，并且所有张量都会被检查以确保尺寸一致。
  • 一个 ...：表示任意数量、任意大小的维度。
  • 一个 str: int 对（实际上是切片），结合了 str 和 int 的行为。（单独的 str 等价于 str: -1。）
  • 一个 str: str 对，此时运行时传入的维度大小将同时绑定到这两个名称，且任何带有其中任一名称的维度都必须具有相同的大小。（有些人喜欢用这种方式为一个维度关联多个名称，以增强文档性。）
  • 一个 str: ... 对，此时与 ... 对应的多个维度将被绑定到 str 指定的名称，并再次检查各参数之间的一致性。
  • None，当与下面的 is_named 一起使用时，表示该维度在 命名张量 的意义上不应具有名称。
  • 一个 None: int 对，结合了 None 和 int 的行为。（单独的 None 等价于 None: -1。）
  • 一个 None: str 对，结合了 None 和 str 的行为。（即该维度不应具有命名，但其大小必须与其他使用该字符串的地方保持一致。）
  • typing.Any：该维度允许任意大小（等同于 -1）。
  • 上述类型的任意元组。例如：TensorType["batch": ..., "length": 10, "channels", -1]。如果只想指定维度的数量，可以使用例如 TensorType[-1, -1, -1] 表示三维张量。
• dtype 参数可以是以下任意一种：
  • torch.float32、torch.float64 等。
  • int、bool、float，这些会被转换为对应的 PyTorch 类型。float 特别地会被解释为 torch.get_default_dtype()，通常为 float32。
• layout 参数可以是 torch.strided 或 torch.sparse_coo，分别用于稠密和稀疏张量。
• details 参数提供了一种传递任意数量附加标志的方式，用于自定义和扩展 torchtyping。默认内置两个标志。torchtyping.is_named 会检查 张量维度的名称，而 torchtyping.is_float 可用于检查是否传入了任意浮点类型。（而不是像 TensorType[torch.float32] 那样只检查特定类型。）有关如何使用自定义 details 来扩展 torchtyping 的讨论，请参阅进一步的文档。
• 可以通过将多个条件组合在一个方括号内来同时检查多项内容。例如：TensorType["batch": ..., "length", "channels", float, is_named]。

torchtyping.patch_typeguard()

torchtyping 与 typeguard 集成，以执行运行时类型检查。应在全局范围内调用 torchtyping.patch_typeguard()，它会修补 typeguard 以检查 TensorType。

此函数可以多次安全调用。（首次调用后便不再执行任何操作。）

• 如果使用 @typeguard.typechecked，则应在使用 @typeguard.typechecked 之前随时调用 torchtyping.patch_typeguard()。例如，可以在每个使用 torchtyping 的文件开头调用。
• 如果使用 typeguard.importhook.install_import_hook，则应在定义要进行检查的函数之前随时调用 torchtyping.patch_typeguard()。例如，可以在安装 typeguard 导入钩子的同时仅调用一次 torchtyping.patch_typeguard()。（钩子和补丁的顺序并不重要。）
• 如果不使用 typeguard，则可以完全省略 torchtyping.patch_typeguard()，仅将 torchtyping 用于文档目的。

pytest --torchtyping-patch-typeguard

torchtyping 提供了一个 pytest 插件，可在测试前自动运行 torchtyping.patch_typeguard()。pytest 会自动发现该插件，您只需传递 --torchtyping-patch-typeguard 标志即可启用它。随后可以像往常一样将包传递给 typeguard，无论是通过 @typeguard.typechecked、typeguard 的导入钩子，还是通过 pytest 标志 --typeguard-packages="your_package_here"。

进一步的文档

请参阅进一步的文档，了解：

• 常见问题解答；
  • 包括与 flake8 和 mypy 的兼容性；
• 如何编写 torchtyping 的自定义扩展；
• 与此主题相关的其他库和资料的资源与链接；
• 更多示例。

torchtyping 快速上手指南

⚠️ 重要提示：原作者强烈建议新项目使用其更新的 jaxtyping 库。它同样支持 PyTorch，不依赖 JAX，且兼容静态类型检查器（如 mypy）。本文档仅针对仍需使用 torchtyping 的旧项目。

torchtyping 允许你为 PyTorch 张量的形状（shape）、数据类型（dtype）、布局及维度名称添加类型注解，并可在运行时进行自动检查，从而消除维度不匹配的 Bug，让代码文档更清晰。

环境准备

在开始之前，请确保满足以下系统要求：

• Python: >= 3.7
• PyTorch: >= 1.7.0
• 可选依赖: typeguard (版本需 < 3.0.0)，用于启用运行时类型检查。

安装步骤

使用 pip 进行安装。国内用户推荐使用清华或阿里镜像源以加速下载。

# 使用默认源
pip install torchtyping

# 或使用国内镜像源（推荐）
pip install torchtyping -i https://pypi.tuna.tsinghua.edu.cn/simple

若需启用运行时检查，请同时安装兼容版本的 typeguard：

pip install "typeguard<3.0.0"

基本使用

1. 基础注解（仅文档用途）

即使不安装 typeguard，你也可以直接使用 TensorType 来替代注释，明确表达张量的维度含义。

import torch
from torchtyping import TensorType

def batch_outer_product(
    x: TensorType["batch", "x_channels"],
    y: TensorType["batch", "y_channels"]
) -> TensorType["batch", "x_channels", "y_channels"]:
    # 代码逻辑保持不变
    return x.unsqueeze(-1) * y.unsqueeze(-2)

2. 启用运行时检查

若要程序在运行时自动验证张量形状是否符合注解，需配合 typeguard 使用。

步骤：

1. 导入并调用 patch_typeguard()（只需调用一次）。
2. 使用 @typechecked 装饰器标记函数。

from torch import rand
from torchtyping import TensorType, patch_typeguard
from typeguard import typechecked

# 必须在使用 @typechecked 之前调用
patch_typeguard()

@typechecked
def add_tensors(
    x: TensorType["batch"],
    y: TensorType["batch"]
) -> TensorType["batch"]:
    return x + y

# 正常情况：形状一致，运行成功
add_tensors(rand(3), rand(3))

# 异常情况：形状不一致，抛出 TypeError
# add_tensors(rand(3), rand(1)) 
# 报错信息示例：TypeError: Dimension 'batch' of inconsistent size. Got both 1 and 3.

3. 高级用法简介

TensorType 支持丰富的参数定义：

• 固定尺寸: 使用整数，如 TensorType[10]。
• 任意尺寸: 使用 -1 或 typing.Any。
• 任意维度数量: 使用 ...，如 TensorType["batch": ..., "channels"]。
• 数据类型: 指定 float, int, torch.float32 等。
• 命名张量: 结合 is_named 标志检查维度名称。

示例：

# 定义一个浮点型、至少两维（第一维为 batch，第二维固定为 10）的张量
def func(x: TensorType["batch": ..., 10, float]) -> None:
    pass