复数求导

jax.grad(lambda z: jnp.sum(jnp.abs(z) ** 2))(z)
jax.grad(lambda z: jnp.sum(jnp.exp(z)), holomorphic=True)(z)
_, pull = jax.vjp(jnp.exp, z1); pull(ct)                 # z1 = [iπ/2], ct = [2-3i]
jax.jvp(jnp.exp, (z1,), (ct,))

(zt.abs() ** 2).sum().backward()                         # PyTorch
torch.autograd.grad(z1t.exp(), z1t, grad_outputs=torch.tensor(ct))
torch.func.jvp(torch.exp, (torch.tensor(z1),), (torch.tensor(ct),))

let z = TracedTensor::from_vec_col_major(vec![2], zd.clone())?;
let m = z.abs()?;                                    // C64 -> F64
let l = (&m * &m)?.reduce_sum(None)?;
ad.grad(&l, &z)
let le = z.exp()?.reduce_sum(None)?;                 // 复数输出
ad.grad(&le, &z)
let y1 = z1.exp()?;
ad.vjp(&y1, &z1, &ct)                                // ct = [2-3i]
ad.vjp(&y1, &z1, &ctc)                               // ctc = [2+3i]，结果再取共轭
ad.jvp(&y1, &z1, &ct)

# JAX
jax   grad(sum|z|^2)            = [0.      -3.141593j, 1.386294+0.j      ]
jax   grad(sum exp z)           = TypeError: grad requires real-valued outputs (output dtype that is a sub-dtype of np.floating), but got complex128. For holomorphic differentiation, pass holomorphic=True. For differentiation of non-holomorphic functions involving complex outputs, use jax.vjp directly.
jax   grad(sum exp z, holomorphic=True) = [0.+1.j, 2.+0.j]
jax   vjp(exp)(2-3i)            = [3.+2.j]  ct*exp z = [3.+2.j]
jax   conj(vjp(exp)(conj(2-3i))) = [-3.-2.j]
jax   jvp(exp)(2-3i)            = [3.+2.j]
# PyTorch
torch backward(sum|z|^2)        = [0.      +3.141593j, 1.386294+0.j      ]
torch backward(sum exp z)       = RuntimeError: grad can be implicitly created only for real scalar outputs but got torch.complex128
torch autograd.grad(sum exp z, grad_outputs=1) = [0.-1.j, 2.+0.j]  conj(exp z) = [0.-1.j, 2.-0.j]
torch vjp(exp)(2-3i)            = [-3.-2.j]  ct*conj(exp z) = [-3.-2.j]
torch jvp(exp)(2-3i)            = [3.+2.j]
# tenferro
abs(z) dtype F64
grad(sum|z|^2)            = [0.000000+3.141593i, 1.386294+0.000000i]
grad(sum exp z)           = [0.000000-1.000000i, 2.000000+0.000000i]  (output dtype C64)
vjp(exp)(2-3i)            = [-3.000000-2.000000i]
conj(vjp(exp)(conj(2-3i))) = [3.000000+2.000000i]
jvp(exp)(2-3i)            = [3.000000+2.000000i]
eager backward(sum|z|^2)  = [0.000000+3.141593i, 1.386294+0.000000i]
eager backward(sum exp z) = [0.000000-1.000000i, 2.000000+0.000000i]

zz = zz - 0.25 * jax.grad(lossj)(zz)            # 直接用 JAX 的梯度
zz = zz - 0.25 * np.conj(jax.grad(lossj)(zz))   # 先取共轭
opt = torch.optim.SGD([zz], lr=0.25)            # PyTorch

let g = ad.grad(&l, &zp)?;    // zp 是 C64 占位符，编译、准备一次，循环里只 run
zz = zz.iter().zip(gv.as_slice::<C>()?).map(|(a, b)| a - b * 0.25).collect();

jax.grad(lambda a: jnp.sum(jnp.linalg.svd(a, compute_uv=False)))(a)
torch.linalg.svdvals(at).sum().backward()
jax.grad(lambda x: jnp.sum(jnp.abs(jnp.einsum("ij,j->i", a, x)) ** 2))(x)

let (_, s, _) = a.svd()?;                 // A 是 C64，s 是 F64
let l = s.reduce_sum(None)?;
ad.grad(&l, &a)

let y = at.dot_general(&xt, dg(&[1], &[0], &[], &[]))?;   // 1. traced dot_general
let y = s.einsum(&[&ea, &ex], "ij,j->i")?;                // 2. eager einsum，再 backward
let vy = tc.einsum(&[va, vx], "ij,j->i")?;                // 3. TraceContext einsum，再 vjp_program

# JAX / PyTorch
L = 4.958876886415667  s = [3.131125, 1.728991, 0.098761]
max|torch - U Vh| = 9.992e-16  max|torch - FD| = 7.769e-10
max|jax - conj(U Vh)| = 0.000e+00  max|jax - FD| = 1.602e+00
torch grad[0, :] = [0.031504+0.800836j, 0.35844 +0.066523j, 0.34888 +0.198334j]
jax   grad[0, :] = [0.031504-0.800836j, 0.35844 -0.066523j, 0.34888 -0.198334j]
max|torch - 2 A^H A x| = 1.986e-15  max|jax - conj(2 A^H A x)| = 8.951e-16
# tenferro
svd s dtype F64
L = 4.958876886415668
max|grad - FD| = 7.769e-10
grad[0, :] = [0.031504+0.800836i, 0.358440+0.066523i, 0.348880+0.198334i]
traced dot_general + ad.grad: max|. - 2 A^H A x| = 1.047e-15
eager einsum + backward: L = 2.719917165456075 max|. - 2 A^H A x| = 1.047e-15
TraceContext einsum + vjp_program: max|. - 2 A^H A x| = 1.047e-15

jax.grad(lambda z: jnp.sum(z * jnp.conj(z)))(z)
jax.grad(lambda z: jnp.real(jnp.sum(z * jnp.conj(z))))(z)
(zt * zt.conj()).sum().real.backward()

let q = (&z * &z.conj()?)?.reduce_sum(None)?;
ad.grad(&q, &z)
z.cast(DType::F64)         // 取实部
z.convert(DType::F64)

# JAX / PyTorch
jax grad: TypeError: grad requires real-valued outputs (output dtype that is a sub-dtype of np.floating), but got complex128. For holomorphic differentiation, pass holomorphic=True. For differentiation of non-holomorphic functions involving complex outputs, use jax.vjp directly.
jax grad(real(sum z conj z)) = [0.      -3.141593j, 1.386294+0.j      ]
torch backward: RuntimeError: grad can be implicitly created only for real scalar outputs but got torch.complex128
torch backward(.real) = [0.      +3.141593j, 1.386294+0.j      ]
# tenferro
sum z conj z: dtype C64 value [2.947854+0.000000i]
grad(sum z conj z)        = [0.000000+3.141593i, 1.386294+0.000000i]
cast(z, F64) = [0.0, 0.6931471805599453]
grad(sum cast(z)^2)       = [0.000000+0.000000i, 1.386294+0.000000i]
convert(z, F64): Err(TracedTensor::convert: unsupported dtype conversion from C64 to F64: checked convert only accepts conversions allowed by dtype promotion; use explicit cast for lossy dtype projection)

jax z -= lr*grad: z = [ 0.999999-6648.51346j , -0.999999-1662.128365j] L = 4.699e+07
jax z -= lr*conj(grad): z = [ 0.999999+1.999998j, -0.999999+0.5j     ] L = 5.684e-12
torch SGD: z = [ 0.999999+1.999998j, -0.999999+0.5j     ] L = 5.684e-12
# tenferro
z -= lr*grad: z = [0.999999+1.999998i, -0.999999+0.500000i] L = 5.684e-12