1D FFTs over multi-dimensional arrays using Arrayfire with Python

Question

Using the pyfftw library, it is easy to do a 1D FFT over a single axis of a multi-dimensional array of shape (M, 2**N) without cycling over the zeroth index explicitly as follows:

afft = pyfftw.interfaces.numpy_fft.fft(a, axis=1).

Switching to arrayfire so that my code can be run on a GPU, I am using the following to do the same thing (with import arrayfire as af):

for idx in af.ParallelRange(M):
    afft = af.fft(a[idx,:])

Although this is executed in parallel it doesn't seem like the best way of doing things and will also make my code horrible to look at.

Is there a way to eliminate the loop using af.fft in a similar way to the pyfftw function, or with an arrayfire version of numpy slicing?