# Example that does a batch of three 2D transformations of size 4 by 5. import torch import pytorch_fft.fft as fft A, zeros = torch.randn( 3 , 4 , 5 ).cuda(), torch.zeros( 3 , 4 , 5 ).cuda() B_real, B_imag = fft.fft2(A_real, A_imag) fft.ifft2(B_real, B_imag) # equals (A, zeros) B_real, B_imag = fft.rfft2(A) # is a truncated version which omits # redundant entries reverse(torch.arange( 0 , 6 )) # outputs [5,4,3,2,1,0] reverse(torch.arange( 0 , 6 ), 2 ) # outputs [4,5,2,3,0,1] expand(B_real) # is equivalent to fft.fft2(A, zeros)[0] expand(B_imag, imag = True ) # is equivalent to fft.fft2(A, zeros)[1]