在 torch.vmap 中高效处理内部张量创建

import torch

poly_batched = torch.tensor([[1, 2, 3, 4], [1, 2, 3, 4]], dtype=torch.float32)

def polycompanion(polynomial):
    # polynomial.shape[-1] 是多项式系数的个数，例如 [a, b, c, d] 代表 ax^3 + bx^2 + cx + d
    # 次数 deg = 系数个数 - 1 - 1 = 系数个数 - 2 (如果最后一个系数是常数项)
    deg = polynomial.shape[-1] - 2

    # 尝试创建伴随矩阵
    companion = torch.zeros((deg + 1, deg + 1), dtype=torch.float32)

    # 填充单位矩阵部分
    companion[1:, :-1] = torch.eye(deg, dtype=torch.float32)

    # 填充最后一列
    # 注意这里 polynomial[:-1] 表示除了最后一个系数以外的所有系数
    # polynomial[-1] 表示最后一个系数
    companion[:, -1] = -1. * polynomial[:-1] / polynomial[-1]
    return companion

# 尝试使用 vmap 向量化
polycompanion_vmap = torch.vmap(polycompanion)

try:
    print(polycompanion_vmap(poly_batched))
except Exception as e:
    print(f"Initial attempt failed: {e}")

def polycompanion_workaround(polynomial, companion_template):
    # 注意：这里的 deg 现在从 companion_template 的形状推断，因为它已经有了批处理维度
    deg = companion_template.shape[-1] - 1 

    # 在传入的 companion_template 上进行就地修改
    companion_template[1:, :-1] = torch.eye(deg, dtype=torch.float32)
    companion_template[:, -1] = -1. * polynomial[:-1] / polynomial[-1]
    return companion_template

polycompanion_vmap_workaround = torch.vmap(polycompanion_workaround)

# 预先创建批处理的 companion 模板
# poly_batched.shape[0] 是批次大小
# poly_batched.shape[-1]-1 是伴随矩阵的行/列维度
companion_init_shape = (poly_batched.shape[0], poly_batched.shape[-1] - 1, poly_batched.shape[-1] - 1)
pre_batched_companion = torch.zeros(companion_init_shape, dtype=torch.float32)

print("--- Workaround Output ---")
print(polycompanion_vmap_workaround(poly_batched, pre_batched_companion))

def polycompanion_optimized(polynomial):
    deg = polynomial.shape[-1] - 2

    # 1. 创建一个基础的非批处理张量来填充单位矩阵部分
    # 这是一个临时的、非批处理的张量
    base_matrix = torch.zeros((deg + 1, deg + 1), dtype=torch.float32)
    base_matrix[1:, :-1] = torch.eye(deg, dtype=torch.float32)

    # 2. 提取 base_matrix 的左侧部分，并进行克隆
    # clone() 创建了一个新的张量，虽然它仍然是非批处理的，
    # 但在 vmap 上下文中，当它与批处理张量拼接时，vmap 会正确处理
    left_part = base_matrix[:, :-1].clone()

    # 3. 计算伴随矩阵的最后一列
    # 这一部分完全从批处理输入 polynomial 派生，因此 vmap 会将其视为批处理张量
    # polynomial[:-1] 是 (deg+1,) 形状
    # polynomial[-1] 是标量
    # 结果是一个 (deg+1,) 形状的张量
    last_column_values = -1. * polynomial[:-1] / polynomial[-1]

    # 4. 扩展最后一列的维度，使其可以与 left_part 进行拼接
    # last_column_values 是 (deg+1,)，我们需要将其变为 (deg+1, 1)
    last_column_reshaped = last_column_values[:, None] 

    # 5. 使用 concatenate 组合左右两部分
    # vmap 会识别 left_part 和 last_column_reshaped，并为它们在批次维度上执行拼接
    final_companion = torch.concatenate([left_part, last_column_reshaped], dim=1)

    return final_companion

polycompanion_vmap_optimized = torch.vmap(polycompanion_optimized)

print("\n--- Optimized Solution Output ---")
print(polycompanion_vmap_optimized(poly_batched))

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