复数梯度下降沿虚轴发散

主线位置：变换主线的「grad」之后，拿梯度去更新复数参数的那一步。

版本：JAX 0.12.0.1（CPU，x64）、PyTorch 2.15.0（单线程）、tenferro-rs 0.8.0（CPU，native，单线程）

flowchart TB
  jx["JAX grad 给 conj 形式"] --> fix["循环里写 conj"]
  fix --> port["原样搬到 tenferro"]
  port --> tg["tenferro grad 已是下降方向"]
  tg --> twice["多取一次共轭"]
  twice --> div["虚部发散"]
  tg --> del["删掉 conj"]
  del --> conv["收敛"]

实验代码：tenferro 一侧是 bin p17_19_complex（release，CpuEngine::serial()），JAX / PyTorch 一侧是 p17_19_complex.py（torch.set_num_threads(1)）。M2 Max、macOS 27、rustc 1.99.0。

现象

最小化 L(z) = Σ|z − w|²，w = [1+2i, −1+0.5i]，从 z = 0 出发，步长 0.25，20 步。JAX 的复数梯度是 ∂L/∂x − i·∂L/∂y，直接减它虚部方向反了，所以 JAX 的写法是先取共轭：

lossj = lambda z: jnp.sum(jnp.abs(z - w) ** 2)
zz = zz - 0.25 * np.conj(jax.grad(lossj)(zz))   # JAX 里正确的写法

jax grad at z=0: [-2.+4.j  2.+1.j]
jax z -= lr*conj(grad): z = [ 0.999999+1.999998j -0.999999+0.5j     ] L = 5.684e-12
jax z -= lr*grad: z = [ 0.999999-6648.51346j  -0.999999-1662.128365j] L = 4.699e+07
torch SGD: z = [ 0.999999+1.999998j -0.999999+0.5j     ]

把这个已经调好的循环搬到 tenferro，conj 一起搬过来：

let dz = (&zp - &w)?.abs()?;
let l = (&dz * &dz)?.reduce_sum(None)?;
let g = ad.grad(&l, &zp)?;
// 每步：zz = zz - 0.25 * conj(g(zz))
zz = zz.iter().zip(&gv).map(|(a, b)| a - b.conj() * 0.25).collect();

grad at z=0: [-2.000000-4.000000i, 2.000000-1.000000i]
z -= lr*conj(grad)  (JAX loop ported): z = [0.999999-6648.513460i, -0.999999-1662.128365i] L = 4.699e7

• 不报错，实部收敛到 Re w，虚部发散，损失从 6.25 涨到 4.7e7。和 JAX 里「忘了取共轭」的结果打印到 6 位小数都相同。
• 起点的梯度两边正好互为共轭：JAX −2+4i，tenferro −2−4i。

出处

三家的导数都对，差在表示：tenferro 和 PyTorch 给 ∂L/∂x + i·∂L/∂y，正是下降方向；JAX 给它的共轭。完整的对照（grad、全纯函数、vjp、jvp 和换算公式）见 复数求导。tenferro 的约定写在 AdContext::grad 的文档里：

    /// For complex scalar outputs, tenferro returns the Hermitian-adjoint
    /// cotangent. To compare seed-`1` scalar gradients with JAX's public
    /// `grad` values, use the complex conjugate of this result.

JAX 循环里的 conj 是为了把 JAX 的表示换成下降方向；tenferro 的梯度已经是下降方向，再取一次共轭，就回到了 JAX 那个虚部反向的量。每步 Im(z − w) 乘 1 + 0.25·2 = 1.5，20 步约 3300 倍。

修法

删掉 conj，和 PyTorch 的 SGD 一样直接减：

zz = zz.iter().zip(&gv).map(|(a, b)| a - b * 0.25).collect();

z -= lr*grad: z = [0.999999+1.999998i, -0.999999+0.500000i] L = 5.684e-12

和 JAX 加 conj、PyTorch SGD 的结果相同。从 JAX 搬过来的代码里，凡是对复数梯度做 conj 的地方（优化器、手写的牛顿步、把梯度和 JAX 参考值比较的测试）都要逐个看：更新参数的删掉 conj，和 JAX 比较的在 tenferro 一侧取共轭再比。

怎么确认没中招

• 在一个已知答案的二次型上跑几步：L = Σ|z − w|² 的梯度在 tenferro 里应当是 2(z − w)，起点 z = 0 时是 −2w。符号对不上，就是多了或少了一个共轭。
• 训练时分别看实部和虚部的残差：只有虚部发散，几乎一定是共轭约定。

小结

• JAX 的复数梯度是 conj(∂L/∂x + i·∂L/∂y)，下降要写 z − lr·conj(g)；tenferro 和 PyTorch 的梯度直接是下降方向。
• JAX 循环原样搬过来会多取一次共轭：不报错，实部收敛，虚部每步放大 1.5 倍。
• 更新参数时删掉 conj；和 JAX 的参考值比较时在 tenferro 一侧取共轭。