implementing conv2d in fourier domain using np.einsum --> ValueError: einstein sum subscripts string contains too many subscripts for operand 0

Question

According to the convolution theorem, convolution operation changes to pointwise multiplication in fourier domain - here I have 'fft_x' of shape (batchsize, height, width, in_channels) which is the fft of input data and similarly 'fft_kernel' of shape (height, width, in_channels, out_channels) which is fft of the kernel after being padded to image size. To get pointwise multiplication of these in efficient way, I was using einsum in the following way -

...
    print(fft_x)
    print(fft_kernel)
    output = 0
    n=int(self.no_of_kernels/2)+1      # n = out_channels here
    for i in range(n):
        output += np.einsum('ijkl,jkl->ijk', fft_x, fft_kernel[i])
    return output 
...

It gives the following output and error-

Tensor("input_11:0", shape=(?, 28, 28, 1), dtype=complex64)
Tensor("fourier__conv2d_11/transpose:0", shape=(28, 28, 1, 17), dtype=complex64)
...
...
ValueError: einstein sum subscripts string contains too many subscripts for operand 0

Could anyone please explain why this error is arising? Thanks in advance for any help.

Answer 1

import tensorflow as tf
output  = tf.einsum('ijkl,jklo->ijko', fft_x, fft_kernel)

I used tf.einsum instead of np.einsum, and it worked. Also as seen in above, I removed the loop and changed the equation in einsum.