![[Bug已解决] max_unpool3d 在 MPS 上对 complex64 输入触发内部 RuntimeError 解决方案](http://pic.xiahunao.cn/yaotu/[Bug已解决] max_unpool3d 在 MPS 上对 complex64 输入触发内部 RuntimeError 解决方案)
[Bug已解决] max_unpool3d 在 MPSApple Silicon上对 complex64 输入触发内部 RuntimeError 解决方案一、现象长什么样你在Apple Silicon 的 MPS 后端Metal Performance Shaders上跑 3D 反池化并且输入是complex64复数张量import torch if torch.backends.mps.is_available(): x torch.randn(2, 3, 4, 4, 4, dtypetorch.complex64, devicemps) # 先 max_pool3d 拿到 indices再 max_unpool3d 还原 out torch.nn.functional.max_unpool3d(x, indices, kernel_size2)结果报torch.nn.functional.max_unpool3d — internal MPS RuntimeError for complex64 input即 pytorch/pytorch#187233。含义MPS 后端的max_unpool3d对 complex64 输入触发了 MPS 内部的 RuntimeError——MPS 内核没有正确处理复数类型的 3D 反池化。这是 MPS 后端「某些算子/ dtype 不支持」类缺口。 本文聚焦max_unpool3d 是什么、为什么 MPS 对 complex64 会内部报错、怎么绕过CPU 计算 / 转实部虚部 / 升级。二、背景max_unpool3d 与 complex64max_unpool3d是max_pool3d的逆操作。max_pool3d 会记录「每个输出位置来自输入的哪个索引indices」max_unpool3d用 indices 把值放回原三维空间其它位置填 0。import torch import torch.nn.functional as F x torch.randn(1, 1, 6, 6, 6) pooled, indices F.max_pool3d(x, kernel_size2, return_indicesTrue) reconstructed F.max_unpool3d(pooled, indices, kernel_size2)complex64复数单精度实部和虚部各 32 位。PyTorch 支持复数张量运算信号处理、FFT 相关。MPS 缺口MPS 后端Apple GPU不是所有算子/dtype 都支持。尤其complex 类型 3D 几何算子这类组合MPS 内核常常没实现或内部断言失败于是报「internal MPS RuntimeError」——这不是你代码错是后端能力缺口。三、为什么 MPS 对 complex64 的 max_unpool3d 会内部报错MPS 的max_unpool3d内核基于 Metal shader大概率只注册了实数类型float32/float16的 shadercomplex64 没有对应 Metal 实现当传入 complex64MPS 的「类型派发」找不到对应内核触发内部MTLComputeCommand错误或断言 →internal MPS RuntimeError这类「复数 3D 池化」本就冷门Apple 的 MPS 团队尚未实现/验证该路径。 与 034 节MPS 上torch.linalg.cross对 bool 报错、056 节MPS 上det/solve未实现同源MPS 后端覆盖不全特定 dtype/算子组合静默缺失或内部崩溃。四、最小可运行复现带守卫下面演示 max_unpool3d 的正常用法CPU 稳以及「MPS complex64」可能崩的示意import torch import torch.nn.functional as F def demo_cpu(): x torch.randn(1, 1, 6, 6, 6, dtypetorch.complex64) pooled, idx F.max_pool3d(x, 2, return_indicesTrue) rec F.max_unpool3d(pooled, idx, 2) print(CPU complex64 max_unpool3d 形状:, rec.shape) def demo_mps(): if not torch.backends.mps.is_available(): print([skip] 无 MPS仅说明) return x torch.randn(1, 1, 6, 6, 6, dtypetorch.complex64, devicemps) pooled, idx F.max_pool3d(x, 2, return_indicesTrue) try: rec F.max_unpool3d(pooled, idx, 2) print(MPS complex64 成功已修复) except RuntimeError as e: if MPS in str(e): print([确认] 命中 MPS complex64 max_unpool3d 内部错误 (#187233):, e) else: print(其它错误:, e) if __name__ __main__: demo_cpu() demo_mps()要点CPU 上 complex64 max_unpool3d 稳定bug 在 MPS 后端。下面给绕过。五、解决方案一在 CPU 上算 max_unpool3d最稳若你的数据不大或该算子非瓶颈把 complex64 的 max_unpool3d 放 CPUimport torch import torch.nn.functional as F def max_unpool3d_cpu(x, idx, ks): if x.device.type mps: x x.to(cpu) # 搬 CPU idx idx.to(cpu) rec F.max_unpool3d(x, idx, ks) return rec # 需要时再搬回 mps # 用法模型主体在 MPS仅该算子回 CPU pooled pooled_mps rec max_unpool3d_cpu(pooled.to(cpu), idx.to(cpu), 2).to(mps)代价一次 CPU/GPU 往返但绝对稳定CPU 全 dtype 支持最完整。六、解决方案二拆实部虚部分别算complex64 的 max_unpool3d 可拆成「实部、虚部各自 max_unpool3d」再拼回——因为 max_unpool3d 是位置还原操作按 indices 放回对复数只是实/虚独立搬运import torch import torch.nn.functional as F def complex_max_unpool3d(x, idx, ks): # 拆实部虚部各自实数分别 unpool再拼 complex64 xr x.real xi x.imag # 注意indices 是对「幅值或实部」算的需保证 pooled 的 indices 适用于实/虚 # 若 pooled 由 complex 的 max_pool3d 得到indices 基于实部比较 # 则实/虚用同一 indices 还原位置一致 ur F.max_unpool3d(xr, idx, ks) ui F.max_unpool3d(xi, idx, ks) return torch.complex(ur, ui) # 若 MPS 对实数 max_unpool3d 支持则可绕过 complex64 的缺口注意这依赖「max_pool3d 对 complex 的 indices 是基于实部比较」——PyTorch 的 complex max_pool 确实按实部取 max与 NumPy 一致。因此实/虚用同一 indices 还原是合法的。如果 MPS 连实数 3D unpool 都支持这招即可完全绕过 complex64 缺口。 更稳的兜底实/虚都在 CPU 算结合方案一。七、解决方案三用 max_pool3d 的 output indices 转回 CPU 路径若你只想要「近似还原」且不在意严格 max_unpool 语义可绕过# 直接在 CPU 用完整 max_unpool3d避免 MPS 内核 rec F.max_unpool3d(pooled.cpu(), indices.cpu(), kernel_size).to(orig_device)八、解决方案四升级并关注修复#187233 是 MPS 后端缺口官方会随版本补全为 max_unpool3d 加 complex64 的 Metal shader或至少优雅报错而非 internal RuntimeError。升级pip install --upgrade torch python -c import torch; print(torch.__version__)判断修复同样 MPS complex64 的max_unpool3d不再 internal RuntimeError或换成清晰的「not implemented」。修复前用「CPU 计算 / 拆实部虚部」绕过。九、排查清单MPS complex64 max_unpool3d报internal MPS RuntimeError且用到max_unpool3d complex64 → 确认 #187233MPS 后端缺口。CPU 正常、MPS 崩溃 → 定位到 MPS 内核不支持 complex64。绕过complex64 的 max_unpool3d 放 CPU 算最稳拆实部虚部分别 unpool 再拼若 MPS 支持实数 3D unpool验证修复后 MPS 上 complex64 不再 internal error通用MPS 冷门 dtype/算子常缺实现遇 internal error 先回 CPU。十、小结torch.nn.functional.max_unpool3d — internal MPS RuntimeError for complex64 input#187233的本质是Apple Silicon 的 MPS 后端max_unpool3d内核只实现了实数类型float32/float16没有 complex64 的 Metal shader传入复数时类型派发失败触发 MPS 内部 RuntimeError。这是 MPS 后端「算子/dtype 覆盖不全」的又一类缺口。 应对CPU 计算最稳complex64 的 max_unpool3d 放 CPUCPU 全 dtype 支持最完整结果再搬回 MPS拆实部虚部max_unpool3d 是「按 indices 位置还原」实/虚独立 unpool 再torch.complex拼回若 MPS 支持实数 3D unpool 即可绕过升级官方补全 complex64 的 MPS max_unpool3d 后修复通用经验MPS 冷门 dtype/算子常缺实现遇internal MPS RuntimeError先回 CPU 验证是否为后端缺口。 记住MPS 不是所有 dtype/算子都支持。complex64 3D 几何算子这类冷门组合MPS 内核常缺失回 CPU 是最可靠的兜底。