Python 如何为tensorflow中处理复张量的自定义op定义梯度？

如果你有一个复数张量，tensorflow已经有了一个运算来获得这些复数的大小（）。现在我想添加一个函数来计算每个数字的角度。这个函数很容易用numpy实现，并且梯度很小

---

在添加基于numpy的自定义tensorflow操作时，这似乎是一个不错的选择。然而，这里的输入是复数。所以我的问题是，是否可以使用相同的方法，但处理复数？直接复制TF.ABS中使用的方法是不可能的，因为它使用C++。

看起来像<>代码> tf。ARG< /Cord>将很快存在：如果你想在那之前做点什么，我推荐你

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

---

太好了，他们正在把它添加到Tensorflow本身中！以下是关于tf.angle的Tensorflow文档链接：