在使用pytorch的`torch.vmap`进行函数向量化时，如果在函数内部创建新的张量（如通过`torch.zeros`），并且该张量的形状不完全由批处理输入直接决定，可能会遇到`batchedtensor`兼容性问题。本文将深入探讨这一挑战，并提供一种优雅的解决方案，通过结合`clone()`和`torch.concatenate`来确保内部创建的张量也能正确地进行批处理，从而实现高效且可读的代码。
理解 torch.vmap 与内部张量创建的挑战
torch.vmap 是 PyTorch 提供的一个强大工具，用于自动向量化函数，将处理单个样本的函数转换为处理一批样本的函数，而无需手动添加批处理维度。它的核心机制是追踪对 BatchedTensor（即 vmap 传入的批处理输入）进行的操作，并确保所有派生张量都正确地保留批处理维度。
然而，当我们在 vmap 内部的函数中，使用 torch.zeros 或 torch.empty 等函数创建一个新的张量时，如果这个张量的形状不是直接从 BatchedTensor 派生而来，或者我们尝试对其进行原地修改，vmap 可能无法正确识别并对其进行批处理。
考虑以下计算多项式伴随矩阵的函数 polycompanion：
import torch
def polycompanion(polynomial):
# polynomial 是一个一维张量，表示多项式系数
# 例如 [a_n, a_{n-1}, ..., a_0]
# 伴随矩阵的维度通常是 (degree, degree)
# 这里的 deg = polynomial.shape[-1] - 2 是根据具体多项式定义来确定伴随矩阵维度的
deg = polynomial.shape[-1] - 2
# 初始化一个零矩阵作为伴随矩阵的基础
companion = torch.zeros((deg + 1, deg + 1))
# 填充伴随矩阵的左下角部分为单位矩阵
companion[1:, :-1] = torch.eye(deg)
# 填充伴随矩阵的最后一列
# 这一列的值依赖于输入的 polynomial
companion[:, -1] = -1. * polynomial[:-1] / polynomial[-1]
return companion
# 准备批处理输入
poly_batched = torch.tensor([[1, 2, 3, 4], [1, 2, 3, 4]], dtype=torch.float32)
# 尝试使用 vmap 向量化函数
polycompanion_vmap = torch.vmap(polycompanion)
# 运行会失败，因为 companion 在函数内部被创建为普通的 Tensor，
# 并且在 vmap 运行时没有被正确地批处理。
try:
print(polycompanion_vmap(poly_batched))
except Exception as e:
print(f"Error when trying to vmap with internal torch.zeros: {e}")上述代码在 torch.vmap 环境下运行时会抛出错误。原因是 companion = torch.zeros((deg+1, deg+1)) 创建了一个普通的 torch.Tensor，而不是一个 BatchedTensor。尽管 deg 是从 polynomial 的形状推导而来，但 torch.zeros 本身并不知道它需要被批处理。随后的原地赋值操作 companion[1:,:-1] = torch.eye(deg) 和 companion[:,-1] = ... 并没有将 companion 转换为 BatchedTensor，导致 vmap 无法正确处理其批处理维度。
一个常见的临时解决方案是，在调用 vmap 之前预先创建并传入一个批处理的零张量：
# 初始的丑陋 workaround
def polycompanion_workaround(polynomial, companion_template):
# deg = companion_template.shape[-1] - 1 # 这里的deg需要根据template重新计算
deg = polynomial.shape[-1] - 2 # 保持与原函数一致
# 使用传入的模板进行原地修改
companion_template[1:,:-1] = torch.eye(deg)
companion_template[:,-1] = -1. * polynomial[:-1] / polynomial[-1]
return companion_template
polycompanion_vmap_workaround = torch.vmap(polycompanion_workaround)
# 预先创建批处理的零张量
# 注意这里的形状需要与期望的批处理输出形状匹配
batch_size = poly_batched.shape[0]
matrix_dim = poly_batched.shape[-1] - 1 # (deg+1)
pre_batched_zeros = torch.zeros(batch_size, matrix_dim, matrix_dim, dtype=torch.float32)
print("Workaround output:")
print(polycompanion_vmap_workaround(poly_batched, pre_batched_zeros))这种方法虽然能工作，但它要求修改函数签名以接受一个额外的模板张量，并且需要在外部手动管理这个模板的批处理维度，这降低了代码的优雅性和可维护性。
推荐解决方案：结合 clone() 和 torch.concatenate
为了在 vmap 内部优雅地创建并填充批处理张量，我们应该避免对非批处理张量进行原地修改。相反，我们可以通过组合已经批处理的张量片段来构建最终的批处理张量。这里的关键在于利用 clone() 操作来“激活”批处理维度，并使用 torch.concatenate 将这些片段组合起来。
以下是改进后的 polycompanion 函数：
import torch
def polycompanion_optimized(polynomial):
deg = polynomial.shape[-1] - 2
# 1. 创建一个非批处理的零矩阵作为基础结构
# 这个 companion 此时仍是普通的 Tensor
companion_base = torch.zeros((deg + 1, deg + 1), dtype=polynomial.dtype)
# 2. 填充单位矩阵部分
# 这一部分是伴随矩阵的左侧部分
identity_part = companion_base[1:, :-1].clone() # 关键：clone() 使得这部分成为 BatchedTensor
identity_part[torch.eye(deg, dtype=torch.bool)] = 1.0 # 填充单位矩阵
# 3. 计算伴随矩阵的最后一列
# polynomial 是 BatchedTensor，所以这个计算结果自然也是 BatchedTensor
last_column = -1. * polynomial[:-1] / polynomial[-1]
# 4. 扩展 last_column 的维度以匹配 concatenate 的要求
# last_column 的形状是 (deg+1,)，需要变成 (deg+1, 1) 才能与 identity_part 拼接
last_column_expanded = last_column[:, None]
# 5. 使用 torch.concatenate 将批处理的片段组合起来
# identity_part 是 (deg+1, deg) 形状的 BatchedTensor
# last_column_expanded 是 (deg+1, 1) 形状的 BatchedTensor
# 沿着 dim=1 拼接，得到 (deg+1, deg+1) 形状的 BatchedTensor
_companion = torch.concatenate([identity_part, last_column_expanded], dim=1)
return _companion
# 准备批处理输入
poly_batched = torch.tensor([[1, 2, 3, 4], [1, 2, 3, 4]], dtype=torch.float32)
# 使用 vmap 向量化优化后的函数
polycompanion_vmap_optimized = torch.vmap(polycompanion_optimized)
print("\nOptimized vmap output:")
print(polycompanion_vmap_optimized(poly_batched))输出结果：
Optimized vmap output:
tensor([[[ 0.0000, 0.0000, -0.2500],
[ 1.0000, 0.0000, -0.5000],
[ 0.0000, 1.0000, -0.7500]],
[[ 0.0000, 0.0000, -0.2500],
[ 1.0000, 0.0000, -0.5000],
[ 0.0000, 1.0000, -0.7500]]])解决方案解析
- companion_base = torch.zeros((deg + 1, deg + 1), dtype=polynomial.dtype): 我们仍然可以创建一个普通的零矩阵作为基础，用于确定形状。重要的是，我们不直接对它进行原地修改并期望它被批处理。
- identity_part = companion_base[1:, :-1].clone(): 这是关键一步。我们首先从 companion_base 中切片出需要填充单位矩阵的部分。然后，通过调用 .clone()，即使 companion_base 本身不是 BatchedTensor，这个 identity_part 在 vmap 内部执行时，会被 vmap 机制识别为一个需要批处理的张量。它会为每个批次创建一个独立的副本。
- identity_part[torch.eye(deg, dtype=torch.bool)] = 1.0: 在 clone() 之后，我们可以在这个 identity_part 上进行原地修改（填充单位矩阵），因为 vmap 已经将其视为一个批处理的张量。
- *`last_column = -1. polynomial[:-1] / polynomial[-1]**: 这一部分直接从polynomial（一个BatchedTensor）计算而来，因此last_column自然也是一个BatchedTensor`。
- last_column_expanded = last_column[:, None]: 使用 None 或 unsqueeze(-1) 扩展 last_column 的维度，使其形状从 (deg+1,) 变为 (deg+1, 1)，以便与 identity_part 沿着 dim=1 进行拼接。
- _companion = torch.concatenate([identity_part, last_column_expanded], dim=1): 最后，我们将两个已经正确批处理的张量片段 identity_part 和 last_column_expanded 沿着正确的维度拼接起来，形成最终的批处理伴随矩阵 _companion。
注意事项与最佳实践
- 理解 vmap 的批处理传播机制：vmap 通过跟踪对 BatchedTensor 的操作来自动传播批处理维度。如果一个张量在函数内部被创建，但其创建过程没有直接依赖于 BatchedTensor 或没有通过 clone() 等方式“激活”批处理，它将不会被 vmap 自动批处理。
- 避免在 vmap 内部对非批处理张量进行原地修改：尽量避免直接创建 torch.zeros 或 torch.empty 然后对其进行原地修改，除非你确定这些操作不会影响 vmap 的批处理行为，或者你正在修改一个已经被 vmap 识别为 BatchedTensor 的张量。
- 利用 clone() 激活批处理：当需要从一个非批处理的张量结构中派生出需要批处理的部分时，使用 .clone() 是一个有效的策略。它会为每个批次创建一个独立的副本，并允许 vmap 追踪这些副本。
- 构建而非修改：在 vmap 内部，更推荐通过组合已批处理的张量片段来构建目标张量，而不是尝试原地修改一个最初未批处理的张量。
- torch.zeros_like 或 torch.empty_like：如果新张量的形状与某个 BatchedTensor 输入完全一致，使用 torch.zeros_like(input_tensor) 或 torch.empty_like(input_tensor) 可以直接创建批处理的零/空张量，这是最简单直接的方法。但在本例中，伴随矩阵的形状与 polynomial 的形状不同，所以不能直接使用 zeros_like。
总结
在 torch.vmap 中进行动态张量创建时，关键在于确保所有需要批处理的中间张量都能够被 vmap 机制正确识别。通过将张量的创建分解为批处理的片段（例如，利用 clone() 激活批处理维度，或直接从 BatchedTensor 派生），然后使用 torch.concatenate 等操作将它们组合起来，我们可以编写出既高效又符合 vmap 设计哲学的代码。这种方法避免了修改函数签名来传入预批处理张量的“丑陋”做法，使得代码更加清晰和易于维护。
