tensorflow中具有复杂指数的自定义渐变

Question

作为练习，我尝试在Tensorflow中构建一个自定义运算符，并根据Tensorflow API操作组成的同一正向操作的Tensorflow的autodiff检查梯度。但是，我的自定义操作符的渐变是不正确的。我的复杂分析似乎不正确，需要复习一下。

import tensorflow as tf
shape = (1, 16)
dtype = tf.complex64
x = tf.cast(tf.complex(tf.random.normal(shape), tf.random.normal(shape)), dtype=dtype)
def fun(x):
phi = x * tf.math.conj(x)
e = tf.exp(1j * phi)
return e
def d_fun(x):
d_phi = x + tf.math.conj(x)
phi = x * tf.math.conj(x)
d_e = 1j * d_phi * tf.exp(1j * phi)
return d_e
@tf.custom_gradient
def tf_custom(x):    
e = fun(x)
def grad(dy):
d_e = d_fun(x)
return dy * d_e
return e, grad
with tf.GradientTape() as g:
g.watch(x)
res = fun(x)

dy_dx = g.gradient(res, x)
with tf.GradientTape() as g:
g.watch(x)
res2 = tf_custom(x)

dy_dx2 = g.gradient(res2, x)
print(tf.reduce_sum(tf.abs(res - res2)).numpy())
print(tf.reduce_sum(tf.abs(dy_dx - dy_dx2)).numpy())

Answer 1

TensorFlow 2不直接计算复变量函数的导数。它似乎使用Wirtinger微积分将复变量函数的导数计算为实部和虚部的函数。你也可以在这里找到解释。