MacBook Air M2上mlx.core.sqrt()处理复数参数崩溃求助

1. 复数参数调用mlx.core.sqrt()崩溃是否为预期情况？ #

这不是预期行为，报错RuntimeError: [metal::Device] Unable to load kernel v_Sqrtcomplex64complex64说明MLX的Metal后端缺少处理complex64类型平方根的预编译内核，属于MLX针对M2架构GPU的兼容性问题。

可尝试以下临时解决方法：

• 切换到CPU计算：用mx.device.cpu()上下文强制在CPU上执行复数平方根运算，示例代码：

  import mlx.core as mx
  w = 1j * mx.array([0.0, 1.0, 2.0])
  with mx.device.cpu():
      z = mx.sqrt(w)
  print(z)
  

• 手动实现复数平方根：基于复数数学定义拆分计算，示例代码：

  import mlx.core as mx
  def complex_sqrt(x):
      abs_x = mx.abs(x)
      real_part = mx.sqrt((abs_x + x.real) / 2)
      imag_part = mx.sign(x.imag) * mx.sqrt((abs_x - x.real) / 2)
      return real_part + 1j * imag_part
  
  w = 1j * mx.array([0.0, 1.0, 2.0])
  z = complex_sqrt(w)
  print(z)
  

2. GPU轨迹捕获失败的下一步处理 #

报错RuntimeError: [metal::start_capture] Failed to start: Capturing is not supported是因为MLX的start_capture功能对M2架构设备存在支持限制，可尝试以下方向：

• 使用Xcode官方GPU调试工具：打开Xcode，通过Product > Profile启动你的Python脚本，选择GPU Capture模板，直接在Xcode中监控Metal指令执行。
• 更新MLX到最新版本：执行pip install --upgrade mlx mlx-core，新版本可能修复内核缺失和GPU捕获的兼容性问题。
• 提交Bug反馈：将设备信息、完整报错日志整理后提交到MLX官方代码仓库，说明M2设备上复数平方根运算的内核加载失败问题，推动官方修复。

内容的提问来源于stack exchange，提问作者David Banas