[Bug已解决] max_unpool3d 在 MPS 上对 complex64 输入触发内部 RuntimeError 解决方案

[Bug已解决] max_unpool3d 在 MPS（Apple 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, dtype=torch.complex64, device="mps") # 先 max_pool3d 拿到 indices，再 max_unpool3d 还原 out = torch.nn.functional.max_unpool3d(x, indices, kernel_size=2)

结果报：

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 与 complex64

max_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_size=2, return_indices=True) reconstructed = F.max_unpool3d(pooled, indices, kernel_size=2)

complex64：复数单精度（实部和虚部各 32 位）。PyTorch 支持复数张量运算（信号处理、FFT 相关）。MPS 缺口：MPS 后端（Apple GPU）不是所有算子/dtype 都支持。尤其complex 类型 + 3D 几何算子这类组合，MPS 内核常常没实现或内部断言失败，于是报「internal MPS RuntimeError」——这不是你代码错，是后端能力缺口。

三、为什么 MPS 对 complex64 的 max_unpool3d 会内部报错

MPS 的max_unpool3d内核（基于 Metal shader）大概率：

1. 只注册了实数类型（float32/float16）的 shader，complex64 没有对应 Metal 实现；
2. 当传入 complex64，MPS 的「类型派发」找不到对应内核，触发内部MTLComputeCommand错误或断言 →internal MPS RuntimeError；
3. 这类「复数 + 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, dtype=torch.complex64) pooled, idx = F.max_pool3d(x, 2, return_indices=True) 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, dtype=torch.complex64, device="mps") pooled, idx = F.max_pool3d(x, 2, return_indices=True) 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 放 CPU：

import 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）

1. 报internal MPS RuntimeError且用到max_unpool3d+ complex64 → 确认 #187233（MPS 后端缺口）。
2. CPU 正常、MPS 崩溃 → 定位到 MPS 内核不支持 complex64。
3. 绕过：complex64 的 max_unpool3d 放 CPU 算（最稳）；
4. 拆实部虚部分别 unpool 再拼（若 MPS 支持实数 3D unpool）；
5. 验证：修复后 MPS 上 complex64 不再 internal error；
6. 通用：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 放 CPU（CPU 全 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 是最可靠的兜底。