Tfjs: Feature request: FFT (and other tf.spectral ops)

[carlthome]

tfjs has a lot of essential math ops already which is truly fantastic! A big missing primitive for me to be able to fully rely on tfjs for building deep learning music apps is a tf.fft/tf.ifft however. Is this in the pipeline? Perhaps the one in Web Audio's AnalyserNode could be exploited somehow?

[Lewuathe]

I'd like to work on this. But one question about this feature is the input/output type. We need to support complex number in order to complete FFT. TensorFlow uses numpy.complex64 as input/output type. But JavaScript/TypeScript does not support such kind of type natively. Though we can create our own data type, it is necessary to consider the cooperability and compatibility with existing ops and data types.

@dsmilkov @nsthorat Do you have any idea about supporting complex number type in TensorFlow.js?

[faroit]

Its not necessarily required to have a complex data type to implement a DFT. See Implementation of dsp.js.

[nsthorat]

One thing that would be great, if you're willing to tackle this, is to understand how TensorFlow implemented the FFT. Ideally, the FFT would mirror what TensorFlow uses internally.

Regarding WebAudio, the browser spectogram is ever so slightly different than TensorFlow, so we could use that for browser-only applications, but ideally we'd try to use what TensorFlow uses.

If you verify that TensorFlow needs complex numbers, can you circle back here and we can think about whether we need to have complex numbers?

[jotaf98]

This feature would be highly appreciated, time series and speech analysis in particular depend on it.

As for the complex data types, as a temporary yet extendable solution, we could have just a complex64 data type that wraps two float arrays (accessed as .real and .imaginary).

In the long run this datatype could support more ops (add, complex multiply, etc). But pragmatically this is not needed on a first implementation, and can be filled in over time. The people who would like to use an FFT now certainly wouldn't mind being able to construct a complex array from a pair of arrays (real and imaginary parts, held using current data types), calling the FFT function with it, and fetching the real and imaginary arrays from the resulting complex64 object. Anything else involving complex numbers can be implemented using real logic.

[nsthorat]

@jotaf98 I like this approach a lot. For now, it might make sense to prototype something like (via .real / .imaginary) that to make sure an FFT in JS or in WebGL is tractable in a browser, then we can talk about merging it in / implementing a proper complex64 type.

[Lewuathe]

@jotaf98 Thank you for the feedback. That totally makes sense. Let me confirm the expected returned value in this case.

TensorFlow

>>> x = [[1,2,3], [4,5,6]]
>>> tf.fft(x)
<tf.Tensor: id=4, shape=(2, 3), dtype=complex64, numpy=
array([[ 6.       +1.1920929e-07j, -1.4999999+8.6602521e-01j,
        -1.4999998-8.6602551e-01j],
       [15.       -5.9604645e-08j, -1.4999998+8.6602545e-01j,
        -1.4999999-8.6602563e-01j]], dtype=complex64)>

So in that case, we can assume TensorFlow.js returns the Tensor with shape=(2, 3, 2), dtype='float32'. Is this assumption right?

x = tf.tensor2d([[1,2,3], [4,5,6]]);
tf.fft(x) // returns the Tensor with the shape (2,3,2)
[[[ 6., +1.1920929e-07j], [-1.4999999,+8.6602521e-01j], [-1.4999998-8.6602551e-01j]],
  [15., -5.9604645e-08j], [-1.4999998,+8.6602545e-01j], [-1.4999999,-8.6602563e-01j]]]

(cc: @nsthorat )

[jotaf98]

Yes, just adding an extra dimension with size 2 (real/imaginary) would be a good solution!

It may break compatibility in the future when it's replaced by a complex type (as in TensorFlow), however I don't think that it is a big deal if this caveat is mentioned in the doc.

[Lewuathe]

@jotaf98 Thank you for the confirmation!

[nsthorat]

tf.fft() is added. There are separate issues for ifft(), rfft(), etc.

[carlthome]

Great news!