Hi there! I’m a TensorFlow contributor working on improving TensorFlow for signal processing (FFTs, digital filters, complex numbers, framing, windowing, STFT, spectrograms, etc.). You should keep an eye on tf.contrib.signal, which just received support for windowing, framing, and overlap-add. Pretty soon (i.e. in a day or so, it’s out for review) I will check in some code for STFT and inverse STFT that has GPU support and gradient support (so you can make an STFT part of learning, not just for input features). Feel free to file bugs on GitHub with feature requests.
tensorflow - Computation using data flow graphs for scalable machine learninggithub.com
By the way, you should update your code to use tf.spectral.rfft instead of tf.fft — it’s 2x faster and uses less memory! :)