PyTorch vmap中动态张量创建的技巧与最佳实践

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}")

# 初始的丑陋 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))

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]]])