PyTorch 复数张量编译支持torch.compile 中的 ComplexTensor 分解原理与实战指南【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorch导读PyTorch 从 2.14 版本起为torch.compile提供了实验性的复数complex-valued张量编译支持其核心思路不是让编译器直接理解复数运算而是通过ComplexTensor子类把复数运算拆解为实部/虚部的实值运算从而复用大量既有的高性能实值算子尤其是矩阵乘法内核。阅读本文后你将掌握如何在torch.compile中通过enable_complex_wrapper开关启用该能力、理解其底层分解机制源码位于torch/_subclasses/complex_tensor/与torch/_functorch/_aot_autograd/并了解view_as_real等操作因别名语义无法支持的限制。torch.compile 的复数支持是什么、何时可用按照官方用户指南 docs/source/user_guide/torch_compiler/torch.compiler_complex_number_support.md 的说明PyTorch 2.14 开始为复数张量的编译提供了**实验性experimental且默认关闭opt-in**的支持。也就是说即使你的模型包含torch.complex64/torch.complex128张量torch.compile默认也不会走复数分解路径必须显式打开配置开关才会生效。该开关在torch/_functorch/config.py中定义enable_complex_wrapper: bool False见 torch/_functorch/config.pyenable_complex_wrapper默认值为False属于典型的“需要显式 opt-in 的实验特性”。开启它最简单的方式是使用配置补丁上下文管理器torch._functorch.config.patch(...)。快速上手在 torch.compile 中启用复数编译官方文档给出的示例完整复现如下注意原文档示例的with语句行尾缺少冒号以下为修正后可运行的版本import torch import torch._functorch.config def some_function(a: torch.Tensor, b: torch.Tensor) - torch.Tensor: c a b d a * b e torch.sin(c) f torch.cos(d) return torch.atan(f / e) a torch.randn((5, 1), dtypetorch.complex64) b torch.randn((5, 1), dtypetorch.complex64) # Enable compilation of complex-valued tensors with torch._functorch.config.patch(enable_complex_wrapperTrue): out torch.compile(some_function)(a, b)逐段拆解这段代码的要点import torch._functorch.config必须先导入该模块才能访问config.patch这个配置上下文管理器配置对象定义于 torch/_functorch/config.py。config.patch(enable_complex_wrapperTrue)在with块内临时将开关置为True块结束后自动恢复原值。这样可以把复数编译能力限定在需要的代码区间内避免影响其他编译路径。输入张量使用复数 dtype示例中a、b均为torch.complex64即每个元素由两个 float32 组成这是触发复数分解路径的直接条件。torch.compile(some_function)(a, b)函数体内包含加法、乘法、sin、cos、除法与atan都是复数分解表中有实现覆盖的常见算子因此可以被整体编译。需要提醒的是该功能仍处于实验阶段且依赖torch._functorch内部实现属于面向高级用户的前瞻性能力接口在未来版本中可能调整。实现原理ComplexTensor 子类与“实虚分离”存储官方文档明确指出该功能通过torch._subclasses.complex_tensor.ComplexTensor子类实现将复数运算分解为实值运算。在仓库中ComplexTensor类定义于 torch/_subclasses/complex_tensor/_core.py注意complex_tensor是一个包目录而非单个.py文件。两个独立连续张量_re 与 _imComplexTensor的核心数据结构非常直白内部保存_re与_im两个相互独立、各自 contiguous的实值张量分别存放实部与虚部而不是像原生复数张量那样把实虚部交错interleaved存放在同一个 storage 中。参见 torch/_subclasses/complex_tensor/_core.py构造函数会对real与imag调用contiguous()保证后续算子拿到连续内存会校验实部与虚部在 shape、device、dtype、pin_memory 上的一致性通过Tensor._make_wrapper_subclass创建子类实例并把外层 dtype 映射为对应的复数类型float32 - complex64等映射表在_ops/common.py的COMPLEX_TO_REAL/REAL_TO_COMPLEX通过__tensor_flatten__/__tensor_unflatten__暴露内部两个张量供 functorch、AOTAutograd 等子系统展开处理。此外_core.py还提供了两个关键转换入口ComplexTensor.from_interleaved(t)把一个交错布局的复数张量拆成ComplexTensoras_interleaved()调用torch.complex(self.re, self.im)把分离布局重新打包回原生复数张量。这两个方法正是“进入/退出编译块时一次性转换开销”的来源。torch_dispatch与算子查表ComplexTensor通过__torch_dispatch__拦截所有作用在其上的算子调用并在分发表中查找对应的分解实现classmethod def __torch_dispatch__(cls, func, types, args(), kwargsNone): from ._ops.common import lookup_complex kwargs {} if kwargs is None else kwargs impl lookup_complex(func, *args, **kwargs) if impl is None: return NotImplemented return impl(*args, **kwargs)见 torch/_subclasses/complex_tensor/_core.pylookup_complex定义于 torch/_subclasses/complex_tensor/_ops/common.py查找顺序为先查具体 overload再查 overloadpacket最后回退到torch._decomp.get_decompositions得到的分解表。也就是说一个算子只要在COMPLEX_OPS_TABLE中注册了实现或者存在可用的 decomposition就能被复数路径处理。算子注册机制四类实现模板common.py提供了一套声明式的注册装饰器是理解“哪些算子支持”的钥匙register_complex(op, impl)直接注册某个算子如aten.real、aten.imag的复数实现register_simple(op)注册“可对实部、虚部各自独立应用同一算子”的算子如slice、flatten、view、mean、sum、clone、permute、transpose等torch/_subclasses/complex_tensor/_ops/aten.py 中的SIMPLE_OPS_LISTregister_binary_nonlinear(op)注册“乘法类”算子aten.mul、aten.mm等实现时按复数乘法公式展开为四个实值运算real a_r*b_r - a_i*b_iimag a_r*b_i a_i*b_rtorch/_subclasses/complex_tensor/_ops/common.py。源码注释特别指出这种展开的累加顺序与原生复数 BLAS 不同在数值结果上可能产生细微差异register_error(op)显式注册为“不支持”并抛出NotImplementedError用于把已知不支持的算子明确标注出来。从源码结构看aten.py约 1108 行是主要算子实现集合而prims.py覆盖torch._prims层面的原始算子common.py中还有WrapComplexMode/ComplexTensorMode两个TorchDispatchMode用于在训练与推理时把普通复数张量临时包装成ComplexTensor。图编译集成AOTAutograd 中的复数分解 Pass光有张量子类还不够torch.compile的图编译管线AOTAutograd需要显式插入一个“图级复数分解”步骤。当enable_complex_wrapperTrue时编译流程会在 torch/_functorch/_aot_autograd/graph_compile.py 触发if config.enable_complex_wrapper: from .complex_decomposition import decompose_complex_in_graph fw_module decompose_complex_in_graph( fw_module, adjusted_flat_args, aot_config.decompositions )也就是说在compiler(fw_module, adjusted_flat_args)真正调用后端编译器如 Inductor之前先对前向图做一次复数→实数的重写。decompose_complex_in_graph 的细节该 pass 实现在 torch/_functorch/_aot_autograd/complex_decomposition.py。模块 docstring 给出了清晰的定位该 pass 在 functionalization 之后运行因此图中的图是纯函数式的无 mutation、无 aliasing。其工作原理是_has_complex(gm, flat_args)遍历 flat args 与图中call_function节点的node.meta[val]只要发现复数 dtype 张量就进入分解流程complex_decomposition.py_assert_no_incorrect_aliasing_mutation(gm)检查图中是否存在对复数张量的原地修改见下文“别名限制”一节wrapper(*args)把复数输入通过_maybe_wrap包装成ComplexTensor在WrapComplexMode上下文中用fx.Interpreter重跑一遍图让ComplexTensor.__torch_dispatch__自然完成逐算子分解最后用_maybe_unwrap把复数输出重新打包回交错布局用make_fx把这次重跑的结果重新跟踪成一个新的GraphModule。模块 docstring 还特别强调图签名输入/输出数量与 dtype保持不变——复数输入在顶部经aten.real/aten.imag解包复数输出在底部经aten.complex重打包。这意味着对图的下游消费者如 Inductor而言看到的始终是常规实值图。收益与代价收益复用实值硬件内核官方文档明确指出的核心收益是分离布局可以轻易使用那些没有复数版本的优化硬件内核尤其体现在矩阵乘法上。复数mm/matmul在底层往往需要专门的复数实现而拆成实部、虚部后四个实值矩阵乘可以直接复用经过深度调优的实值 BLAS / cuBLAS 内核。这也是“分解式支持”而非“为编译器新增复数代码生成”的根本动机。代价进入/退出编译块的一次性转换同样来自官方文档在进入与退出torch.compile块时需要把交错布局的张量转换为分离的两张张量或反向这是一次性one-time转换成本。对于计算密集、图上算子多的场景这笔开销通常被编译带来的收益摊薄但对于轻量函数或极短生命周期张量转换开销可能相对明显。另外从源码可以推断一个次要代价由于ComplexTensor.__new__强制real.contiguous()/imag.contiguous()见 torch/_subclasses/complex_tensor/_core.py非连续复数张量进入编译区会被显式拷贝为连续布局这可能带来额外的内存与拷贝成本也是文档中“无法保持别名语义”这一限制的根源之一。已知限制别名语义与交错布局算子官方文档的 Limitations 一节指出这种方案最大的限制在于无法为某些算子维持别名aliasing语义尤其是那些本质依赖交错布局的算子。文档点名了两个最典型的例子torch.view_as_realtorch.view_as_complex原因很直观view_as_real/view_as_complex直接操作底层 storage 的内存布局——前者把复数张量“看作”最后一维大小为 2 的实值张量后者反之。在ComplexTensor的实虚分离表示下这两个算子的“视图”语义天然无法成立。另一个常见场景是在torch.compile内部原地修改复数输入张量由于输入在分解时被拆解且转为连续修改无法回写到原始张量。源码侧与之呼应的是一道显式检查complex_decomposition.py中的_assert_no_incorrect_aliasing_mutation会在编译期扫描图中所有aten.copy_节点一旦发现目标 dtype 为复数就抛出运行时错误RuntimeError: Mutating a complex tensor in place is not allowed in a torch.compile region.见 torch/_functorch/_aot_autograd/complex_decomposition.py一个重要的“宽慰”细节文档同时给出了一个容易被忽略的实践要点这些算子在用户代码中出现未必意味着编译失败。因为编译器图优化如算子融合可能把这些操作吸收/融合掉导致它们根本不会出现在最终被编译的图中此时函数依然能成功编译。换言之“源码里用了view_as_complex”与“编译图里有view_as_complex”是两回事实际能否编译取决于融合结果。测试与验证如何确认算子支持情况仓库为复数支持维护了一套专门的算子覆盖测试test/complex_tensor/test_complex_tensor.pyOwner 标记为module: complex。测试文件通过implemented_op_db、force_test_op_db等机制批量遍历已注册算子并用SKIPS字典显式跳过已知问题项。值得注意的跳过项包括aten.empty_like/aten.randn_like非确定性输出aten.any/aten.all/aten.allclose等算子的分布式变体没有注册分片策略aten.view_as_real的分布式变体无标量支持——这与文档中点名的view_as_real限制相互印证。如果你在自己模型中发现某个复数算子无法编译可以先核对COMPLEX_OPS_TABLE与DECOMPOSITIONS是否覆盖了该算子再决定是否向官方反馈。反馈不支持的操作官方文档建议如果你需要支持的复数算子目前无法编译可以先查阅官方 issue 列表中带module: complex与module: functorch两个标签的开放问题该列表由 PyTorch 官方仓库维护如果没有已存在的 issue则为该算子新建一个 issue 并提供最小复现。考虑到文档与实现均标注为“实验性”复数编译的支持面仍在持续扩充中反馈是推动覆盖范围扩大的主要途径。小结PyTorch 2.14 起的复数编译支持走了一条务实的技术路线以ComplexTensor子类torch/_subclasses/complex_tensor/_core.py为载体、以“实虚分离 逐算子查表分解”为手段在 AOTAutograd 图编译阶段通过decompose_complex_in_graphtorch/_functorch/_aot_autograd/complex_decomposition.py把复数图整体改写为实值图从而复用成熟的实值内核。开启方式只需一行配置torch._functorch.config.patch(enable_complex_wrapperTrue)。使用前请务必评估两个前提进入/退出编译块的一次性转换开销以及view_as_real/view_as_complex/ 原地修改复数输入等别名敏感操作可能导致的编译失败或显式报错。【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorch创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 PyTorch 复数张量编译支持：torch.compile 中的 ComplexTensor 分解原理与实战指南
PyTorch 复数张量编译支持torch.compile 中的 ComplexTensor 分解原理与实战指南【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorch导读PyTorch 从 2.14 版本起为torch.compile提供了实验性的复数complex-valued张量编译支持其核心思路不是让编译器直接理解复数运算而是通过ComplexTensor子类把复数运算拆解为实部/虚部的实值运算从而复用大量既有的高性能实值算子尤其是矩阵乘法内核。阅读本文后你将掌握如何在torch.compile中通过enable_complex_wrapper开关启用该能力、理解其底层分解机制源码位于torch/_subclasses/complex_tensor/与torch/_functorch/_aot_autograd/并了解view_as_real等操作因别名语义无法支持的限制。torch.compile 的复数支持是什么、何时可用按照官方用户指南 docs/source/user_guide/torch_compiler/torch.compiler_complex_number_support.md 的说明PyTorch 2.14 开始为复数张量的编译提供了**实验性experimental且默认关闭opt-in**的支持。也就是说即使你的模型包含torch.complex64/torch.complex128张量torch.compile默认也不会走复数分解路径必须显式打开配置开关才会生效。该开关在torch/_functorch/config.py中定义enable_complex_wrapper: bool False见 torch/_functorch/config.pyenable_complex_wrapper默认值为False属于典型的“需要显式 opt-in 的实验特性”。开启它最简单的方式是使用配置补丁上下文管理器torch._functorch.config.patch(...)。快速上手在 torch.compile 中启用复数编译官方文档给出的示例完整复现如下注意原文档示例的with语句行尾缺少冒号以下为修正后可运行的版本import torch import torch._functorch.config def some_function(a: torch.Tensor, b: torch.Tensor) - torch.Tensor: c a b d a * b e torch.sin(c) f torch.cos(d) return torch.atan(f / e) a torch.randn((5, 1), dtypetorch.complex64) b torch.randn((5, 1), dtypetorch.complex64) # Enable compilation of complex-valued tensors with torch._functorch.config.patch(enable_complex_wrapperTrue): out torch.compile(some_function)(a, b)逐段拆解这段代码的要点import torch._functorch.config必须先导入该模块才能访问config.patch这个配置上下文管理器配置对象定义于 torch/_functorch/config.py。config.patch(enable_complex_wrapperTrue)在with块内临时将开关置为True块结束后自动恢复原值。这样可以把复数编译能力限定在需要的代码区间内避免影响其他编译路径。输入张量使用复数 dtype示例中a、b均为torch.complex64即每个元素由两个 float32 组成这是触发复数分解路径的直接条件。torch.compile(some_function)(a, b)函数体内包含加法、乘法、sin、cos、除法与atan都是复数分解表中有实现覆盖的常见算子因此可以被整体编译。需要提醒的是该功能仍处于实验阶段且依赖torch._functorch内部实现属于面向高级用户的前瞻性能力接口在未来版本中可能调整。实现原理ComplexTensor 子类与“实虚分离”存储官方文档明确指出该功能通过torch._subclasses.complex_tensor.ComplexTensor子类实现将复数运算分解为实值运算。在仓库中ComplexTensor类定义于 torch/_subclasses/complex_tensor/_core.py注意complex_tensor是一个包目录而非单个.py文件。两个独立连续张量_re 与 _imComplexTensor的核心数据结构非常直白内部保存_re与_im两个相互独立、各自 contiguous的实值张量分别存放实部与虚部而不是像原生复数张量那样把实虚部交错interleaved存放在同一个 storage 中。参见 torch/_subclasses/complex_tensor/_core.py构造函数会对real与imag调用contiguous()保证后续算子拿到连续内存会校验实部与虚部在 shape、device、dtype、pin_memory 上的一致性通过Tensor._make_wrapper_subclass创建子类实例并把外层 dtype 映射为对应的复数类型float32 - complex64等映射表在_ops/common.py的COMPLEX_TO_REAL/REAL_TO_COMPLEX通过__tensor_flatten__/__tensor_unflatten__暴露内部两个张量供 functorch、AOTAutograd 等子系统展开处理。此外_core.py还提供了两个关键转换入口ComplexTensor.from_interleaved(t)把一个交错布局的复数张量拆成ComplexTensoras_interleaved()调用torch.complex(self.re, self.im)把分离布局重新打包回原生复数张量。这两个方法正是“进入/退出编译块时一次性转换开销”的来源。torch_dispatch与算子查表ComplexTensor通过__torch_dispatch__拦截所有作用在其上的算子调用并在分发表中查找对应的分解实现classmethod def __torch_dispatch__(cls, func, types, args(), kwargsNone): from ._ops.common import lookup_complex kwargs {} if kwargs is None else kwargs impl lookup_complex(func, *args, **kwargs) if impl is None: return NotImplemented return impl(*args, **kwargs)见 torch/_subclasses/complex_tensor/_core.pylookup_complex定义于 torch/_subclasses/complex_tensor/_ops/common.py查找顺序为先查具体 overload再查 overloadpacket最后回退到torch._decomp.get_decompositions得到的分解表。也就是说一个算子只要在COMPLEX_OPS_TABLE中注册了实现或者存在可用的 decomposition就能被复数路径处理。算子注册机制四类实现模板common.py提供了一套声明式的注册装饰器是理解“哪些算子支持”的钥匙register_complex(op, impl)直接注册某个算子如aten.real、aten.imag的复数实现register_simple(op)注册“可对实部、虚部各自独立应用同一算子”的算子如slice、flatten、view、mean、sum、clone、permute、transpose等torch/_subclasses/complex_tensor/_ops/aten.py 中的SIMPLE_OPS_LISTregister_binary_nonlinear(op)注册“乘法类”算子aten.mul、aten.mm等实现时按复数乘法公式展开为四个实值运算real a_r*b_r - a_i*b_iimag a_r*b_i a_i*b_rtorch/_subclasses/complex_tensor/_ops/common.py。源码注释特别指出这种展开的累加顺序与原生复数 BLAS 不同在数值结果上可能产生细微差异register_error(op)显式注册为“不支持”并抛出NotImplementedError用于把已知不支持的算子明确标注出来。从源码结构看aten.py约 1108 行是主要算子实现集合而prims.py覆盖torch._prims层面的原始算子common.py中还有WrapComplexMode/ComplexTensorMode两个TorchDispatchMode用于在训练与推理时把普通复数张量临时包装成ComplexTensor。图编译集成AOTAutograd 中的复数分解 Pass光有张量子类还不够torch.compile的图编译管线AOTAutograd需要显式插入一个“图级复数分解”步骤。当enable_complex_wrapperTrue时编译流程会在 torch/_functorch/_aot_autograd/graph_compile.py 触发if config.enable_complex_wrapper: from .complex_decomposition import decompose_complex_in_graph fw_module decompose_complex_in_graph( fw_module, adjusted_flat_args, aot_config.decompositions )也就是说在compiler(fw_module, adjusted_flat_args)真正调用后端编译器如 Inductor之前先对前向图做一次复数→实数的重写。decompose_complex_in_graph 的细节该 pass 实现在 torch/_functorch/_aot_autograd/complex_decomposition.py。模块 docstring 给出了清晰的定位该 pass 在 functionalization 之后运行因此图中的图是纯函数式的无 mutation、无 aliasing。其工作原理是_has_complex(gm, flat_args)遍历 flat args 与图中call_function节点的node.meta[val]只要发现复数 dtype 张量就进入分解流程complex_decomposition.py_assert_no_incorrect_aliasing_mutation(gm)检查图中是否存在对复数张量的原地修改见下文“别名限制”一节wrapper(*args)把复数输入通过_maybe_wrap包装成ComplexTensor在WrapComplexMode上下文中用fx.Interpreter重跑一遍图让ComplexTensor.__torch_dispatch__自然完成逐算子分解最后用_maybe_unwrap把复数输出重新打包回交错布局用make_fx把这次重跑的结果重新跟踪成一个新的GraphModule。模块 docstring 还特别强调图签名输入/输出数量与 dtype保持不变——复数输入在顶部经aten.real/aten.imag解包复数输出在底部经aten.complex重打包。这意味着对图的下游消费者如 Inductor而言看到的始终是常规实值图。收益与代价收益复用实值硬件内核官方文档明确指出的核心收益是分离布局可以轻易使用那些没有复数版本的优化硬件内核尤其体现在矩阵乘法上。复数mm/matmul在底层往往需要专门的复数实现而拆成实部、虚部后四个实值矩阵乘可以直接复用经过深度调优的实值 BLAS / cuBLAS 内核。这也是“分解式支持”而非“为编译器新增复数代码生成”的根本动机。代价进入/退出编译块的一次性转换同样来自官方文档在进入与退出torch.compile块时需要把交错布局的张量转换为分离的两张张量或反向这是一次性one-time转换成本。对于计算密集、图上算子多的场景这笔开销通常被编译带来的收益摊薄但对于轻量函数或极短生命周期张量转换开销可能相对明显。另外从源码可以推断一个次要代价由于ComplexTensor.__new__强制real.contiguous()/imag.contiguous()见 torch/_subclasses/complex_tensor/_core.py非连续复数张量进入编译区会被显式拷贝为连续布局这可能带来额外的内存与拷贝成本也是文档中“无法保持别名语义”这一限制的根源之一。已知限制别名语义与交错布局算子官方文档的 Limitations 一节指出这种方案最大的限制在于无法为某些算子维持别名aliasing语义尤其是那些本质依赖交错布局的算子。文档点名了两个最典型的例子torch.view_as_realtorch.view_as_complex原因很直观view_as_real/view_as_complex直接操作底层 storage 的内存布局——前者把复数张量“看作”最后一维大小为 2 的实值张量后者反之。在ComplexTensor的实虚分离表示下这两个算子的“视图”语义天然无法成立。另一个常见场景是在torch.compile内部原地修改复数输入张量由于输入在分解时被拆解且转为连续修改无法回写到原始张量。源码侧与之呼应的是一道显式检查complex_decomposition.py中的_assert_no_incorrect_aliasing_mutation会在编译期扫描图中所有aten.copy_节点一旦发现目标 dtype 为复数就抛出运行时错误RuntimeError: Mutating a complex tensor in place is not allowed in a torch.compile region.见 torch/_functorch/_aot_autograd/complex_decomposition.py一个重要的“宽慰”细节文档同时给出了一个容易被忽略的实践要点这些算子在用户代码中出现未必意味着编译失败。因为编译器图优化如算子融合可能把这些操作吸收/融合掉导致它们根本不会出现在最终被编译的图中此时函数依然能成功编译。换言之“源码里用了view_as_complex”与“编译图里有view_as_complex”是两回事实际能否编译取决于融合结果。测试与验证如何确认算子支持情况仓库为复数支持维护了一套专门的算子覆盖测试test/complex_tensor/test_complex_tensor.pyOwner 标记为module: complex。测试文件通过implemented_op_db、force_test_op_db等机制批量遍历已注册算子并用SKIPS字典显式跳过已知问题项。值得注意的跳过项包括aten.empty_like/aten.randn_like非确定性输出aten.any/aten.all/aten.allclose等算子的分布式变体没有注册分片策略aten.view_as_real的分布式变体无标量支持——这与文档中点名的view_as_real限制相互印证。如果你在自己模型中发现某个复数算子无法编译可以先核对COMPLEX_OPS_TABLE与DECOMPOSITIONS是否覆盖了该算子再决定是否向官方反馈。反馈不支持的操作官方文档建议如果你需要支持的复数算子目前无法编译可以先查阅官方 issue 列表中带module: complex与module: functorch两个标签的开放问题该列表由 PyTorch 官方仓库维护如果没有已存在的 issue则为该算子新建一个 issue 并提供最小复现。考虑到文档与实现均标注为“实验性”复数编译的支持面仍在持续扩充中反馈是推动覆盖范围扩大的主要途径。小结PyTorch 2.14 起的复数编译支持走了一条务实的技术路线以ComplexTensor子类torch/_subclasses/complex_tensor/_core.py为载体、以“实虚分离 逐算子查表分解”为手段在 AOTAutograd 图编译阶段通过decompose_complex_in_graphtorch/_functorch/_aot_autograd/complex_decomposition.py把复数图整体改写为实值图从而复用成熟的实值内核。开启方式只需一行配置torch._functorch.config.patch(enable_complex_wrapperTrue)。使用前请务必评估两个前提进入/退出编译块的一次性转换开销以及view_as_real/view_as_complex/ 原地修改复数输入等别名敏感操作可能导致的编译失败或显式报错。【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorch创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考