Perform tf.signal.fft2d in the middle of the tensor

Question

I have a tf.Tensor of, for example, shape (31, 6, 6, 3).

I want to perform tf.signal.fft2d on the shapes 6, 6 so, in other words, in the middle. However, the description says:

Computes the 2-dimensional discrete Fourier transform over the inner-most 2 dimensions of input

I could do it with a for loop but I fear it might be very ineffective. Is there a fastest way?

The result must have the same output shape of course.

Answer 1

Thanks to this I implemented this solution using tf.transpose:

in_pad = tf.transpose(in_pad, perm=[0, 3, 1, 2])
out = tf.signal.fft2d(tf.cast(in_pad, tf.complex64))
out = tf.transpose(out, perm=[0, 2, 3, 1])