如何为在 tensorflow 中处理复杂张量的自定义操作定义渐变？

Question

如果您有一个复数张量，则 tensorflow 已经有一个操作来获取这些复数的大小 (tf.abs)。现在我想添加一个计算每个数字的角度的函数。这个函数使用 numpy 很容易实现，梯度为known。

在基于 numpy 添加自定义 tensorflow 操作时，this 似乎是要走的路。然而，这里的输入是复数。所以我的问题是，是否可以使用相同的方法但使用复数？直接复制 tf.abs 中使用的方法是不可能的，因为它使用的是 c++。

Answer 1

看起来tf.arg 很快就会存在：https://github.com/tensorflow/tensorflow/pull/10643。如果你在那之前想要一些东西，我会推荐

def arg(z):
  return tf.atan2(tf.imag(z), tf.real(z))