python - 如何为在 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))