fft#
-
nvmath.
fft. fft( - operand,
- axes=None,
- direction=None,
- options=None,
- execution=None,
- prolog=None,
- epilog=None,
- stream=None,
Perform an N-D complex-to-complex (C2C) FFT on the provided complex operand.
- Parameters:
operand – A tensor (ndarray-like object). The currently supported types are
numpy.ndarray
,cupy.ndarray
, andtorch.Tensor
.axes – The dimensions along which the FFT is performed.
axes[-1]
is the ‘last transformed’ axis for rffts. Currently, it is required that the axes are contiguous and include the first or the last dimension. Only up to 3D FFTs are supported.options – Specify options for the FFT as a
FFTOptions
object. Alternatively, adict
containing the parameters for theFFTOptions
constructor can also be provided. If not specified, the value will be set to the default-constructedFFTOptions
object.prolog – Provide device-callable function in LTO-IR format to use as load-callback as an object of type
DeviceCallable
. Alternatively, adict
containing the parameters for theDeviceCallable
constructor can also be provided. The default is no prolog. Currently, callbacks are supported only with CUDA execution.epilog – Provide device-callable function in LTO-IR format to use as store-callback as an object of type
DeviceCallable
. Alternatively, adict
containing the parameters for theDeviceCallable
constructor can also be provided. The default is no epilog. Currently, callbacks are supported only with CUDA execution.stream – Provide the CUDA stream to use for executing the operation. Acceptable inputs include
cudaStream_t
(as Pythonint
),cupy.cuda.Stream
, andtorch.cuda.Stream
. If a stream is not provided, the current stream from the operand package will be used.
- Returns:
A transformed operand that retains the same data type and shape as the input. It remains on the same device and uses the same package as the input operand.
Examples
>>> import cupy as cp >>> import nvmath
Create a 3-D complex128 ndarray on the GPU:
>>> shape = 256, 256, 256 >>> a = cp.random.rand(*shape, dtype=cp.float64) + 1j * cp.random.rand(*shape, dtype=cp.float64)
Perform a 3-D C2C FFT using
fft()
. The resultr
is also a CuPy complex128 ndarray:>>> r = nvmath.fft.fft(a)
User may also perform FFT along a subset of dimensions, e.g, 2-D C2C FFT along the first two dimensions, batched along the last dimension:
>>> axes = 0, 1 >>> r = nvmath.fft.fft(a, axes=axes)
For C2C type FFT operation, the output can be directly computed inplace thus overwriting the input operand. This can be specified using options to the FFT:
>>> o = nvmath.fft.FFTOptions(inplace=True) >>> r = nvmath.fft.fft(a, options=o) >>> r is a
See
FFTOptions
for the complete list of available options.The package current stream is used by default, but a stream can be explicitly provided to the FFT operation. This can be done if the FFT operand is computed on a different stream, for example:
>>> s = cp.cuda.Stream() >>> with s: ... a = cp.random.rand(*shape) + 1j * cp.random.rand(*shape) >>> nvmath.fft.fft(a, stream=s)
The operation above runs on stream
s
and is ordered with respect to the input computation.Create a NumPy ndarray on the CPU.
>>> import numpy as np >>> b = np.random.rand(*shape) + 1j * np.random.rand(*shape)
Provide the NumPy ndarray to
fft()
, with the result also being a NumPy ndarray:>>> r = nvmath.fft.fft(b)
Notes
This function only takes complex operand for C2C transformation. If the user wishes to perform full FFT transformation on real input, please cast the input to the corresponding complex data type.
This function is a convenience wrapper around
FFT
and and is specifically meant for single use. The same computation can be performed with the stateful API using the defaultdirection
argument inFFT.execute()
.
Further examples can be found in the nvmath/examples/fft directory.