Wie definiere ich einen Farbverlauf für einen benutzerdefinierten Op, der an komplexen Tensoren im Tensorflow arbeitet?

Question

Wenn Sie einen Tensor komplexer Zahlen haben, hat Tensorflow bereits einen Op, um die Größen dieser komplexen Zahlen zu erhalten (tf.abs). Jetzt möchte ich eine Funktion hinzufügen, die den Winkel jeder Zahl berechnet. Diese Funktion lässt sich einfach mit numpy implementieren und der Gradient ist known.

Wenn es darum geht, benutzerdefinierte Tensorflow-Ops basierend auf numpy hinzuzufügen, scheint this der Weg zu sein. Jedoch sind hier die Eingaben komplexe Zahlen. Meine Frage ist also, ist es möglich, denselben Ansatz zu verwenden, aber mit komplexen Zahlen zu arbeiten? Das direkte Kopieren des in tf.abs verwendeten Ansatzes ist nicht möglich, da hier C++ verwendet wird.

Answer 1

Sieht aus wie tf.arg wird in Kürze existieren: https://github.com/tensorflow/tensorflow/pull/10643. Wenn Sie etwas vorher wollen, würde ich empfehlen

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