nvmath.fft.fft — NVIDIA nvmath-python (original) (raw)

nvmath.fft.fft(

x,

*,

axes=None,

direction=None,

options=None,

execution=None,

prolog=None,

epilog=None,

stream=None,

check_dtype=None,

)[source]#

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, and torch.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, a dict containing the parameters for the FFTOptions constructor can also be provided. If not specified, the value will be set to the default-constructed FFTOptions object.
• prolog – Provide device-callable function in LTO-IR format to use as load-callback as an object of type DeviceCallable. Alternatively, a dict containing the parameters for the DeviceCallable 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, a dict containing the parameters for the DeviceCallable 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 Python int), cupy.cuda.Stream, and torch.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 result r 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 True

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) r = 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 default direction argument in FFT.execute().

Further examples can be found in the nvmath/examples/fft directory.