Pytorch version of my numpy code not running vectorized

Hi, I’m currently converting my numpy vectorized code running a fft2 to pytorch, but somehow the result is different. Am I doing it wrong?

numpy version:

def apply_filter_vector(_filter, arr):
    tc= arr.copy()
    temp = apply_filter(tc, _filter[:, :, None])
    return temp 

def apply_filter(arr, _filter, axes=(0,1)):
    return ifft2(
        np.multiply(ifftshift(1 - _filter), fft2(arr, axes=axes)),
        axes=axes,
    ).real

pytorch version :

def apply_filter_vector_pytorch(_filter, arr, cuda=False):
    tc= arr.copy()
    
    if cuda == True:
        temp = apply_filter_pytorch(tc, _filter[:, :, None], cuda=True)        
    else:
        temp = apply_filter_pytorch(tc, _filter[:, :, None], cuda=False) 
    return temp

def apply_filter_pytorch(arr, _filter, axes=(0,1), cuda=False):
    _filter = torch.from_numpy(_filter)
    arr = torch.from_numpy(arr)
    if cuda == True:
        _filter=_filter.cuda()
        arr=arr.cuda()
    result = torch.fft.ifft2(torch.multiply(torch.fft.ifftshift(1.0-_filter),torch.fft.fft2(arr, dim =axes)))
    if cuda == True:
        result= result.cpu()
    result= result.numpy().real
    return result

the data is a 3D array (2D square images with z axis variation in time).

Could you post the shapes for all inputs, as the code fails in a broadcast:

    np.multiply(np.fft.ifftshift(1 - _filter), np.fft.fft2(arr, axes=axes))

ValueError: operands could not be broadcast together with shapes (24,24,1) (100,100,100) 

using:

filt = np.random.randn(24, 24)
arr = np.random.randn(100, 100, 100)

out_np = apply_filter_vector(filt, arr)

Also, how large is the abs().max() error you are seeing?

hi @ptrblck , Thanks, I already solved it, turns out I only had to add the dim to the ifft2!

torch.fft.ifft2(
        torch.multiply(
            torch.fft.ifftshift(1.0 - _filter), torch.fft.fft2(arr, dim=axes)
        ),
        dim=axes,
    )