If you have a tensor of complex numbers, tensorflow already has an op to get the magnitudes of those complex numbers (tf.abs). Now I want to add a function that calculates the angle of each number. This function is easy to implement using numpy and the gradient is known.
When it comes to adding custom tensorflow ops based on numpy, this seems to be the way to go. However here the input are complex numbers. So my question is, is it possible to use the same approach but work with complex numbers? Just directly copying the approach used in tf.abs is not possible, since it uses c++.
答案 0 :(得分：2)
看起来tf.arg很快就会存在：https://github.com/tensorflow/tensorflow/pull/10643。如果你之前想要的东西，我建议
def arg(z):
return tf.atan2(tf.imag(z), tf.real(z))