fft#
-
nvmath.
fft. fft( - x,
- *,
- axes=None,
- direction=None,
- 选项=None,
- 执行=None,
- 前言=None,
- 后记=None,
- 流=None,
- 检查数据类型=None,
fft(operand, axes=None, direction=None, options=None, execution=None, prolog=None, epilog=None, stream=None)
对提供的复数操作数执行 N 维复数到复数 (C2C) FFT。
- 参数:
operand – 张量(类似 ndarray 的对象)。目前支持的类型为
numpy.ndarray
,cupy.ndarray
, 和torch.Tensor
。axes – 执行 FFT 的维度。
axes[-1]
是 rfft 的“最后转换”轴。目前,要求轴是连续的,并且包含第一个或最后一个维度。仅支持高达 3D FFT。options – 将 FFT 的选项指定为
FFTOptions
对象。或者,也可以提供包含FFTOptions
构造函数参数的dict
。如果未指定,则该值将设置为默认构造的FFTOptions
对象。prolog – 以 LTO-IR 格式提供设备可调用函数,以用作加载回调,作为
DeviceCallable
类型的对象。或者,也可以提供包含DeviceCallable
构造函数参数的dict
。默认情况下没有 prolog。目前,回调仅在 CUDA 执行中受支持。epilog – 以 LTO-IR 格式提供设备可调用函数,以用作存储回调,作为
DeviceCallable
类型的对象。或者,也可以提供包含DeviceCallable
构造函数参数的dict
。默认情况下没有 epilog。目前,回调仅在 CUDA 执行中受支持。stream – 提供用于执行操作的 CUDA 流。可接受的输入包括
cudaStream_t
(作为 Pythonint
),cupy.cuda.Stream
, 和torch.cuda.Stream
。如果未提供流,将使用操作数包中的当前流。
- 返回:
转换后的操作数,它保留与输入相同的数据类型和形状。它保留在同一设备上,并使用与输入操作数相同的包。
示例
>>> import cupy as cp >>> import nvmath
在 GPU 上创建一个 3D complex128 ndarray
>>> shape = 256, 256, 256 >>> a = cp.random.rand(*shape, dtype=cp.float64) + 1j * cp.random.rand( ... *shape, dtype=cp.float64 ... )
使用
fft()
执行 3D C2C FFT。结果r
也是一个 CuPy complex128 ndarray>>> r = nvmath.fft.fft(a)
用户还可以沿维度子集执行 FFT,例如,沿前两个维度执行 2D C2C FFT,沿最后一个维度进行批处理
>>> axes = 0, 1 >>> r = nvmath.fft.fft(a, axes=axes)
对于 C2C 类型 FFT 操作,可以直接就地计算输出,从而覆盖输入操作数。这可以使用 FFT 的选项来指定
>>> o = nvmath.fft.FFTOptions(inplace=True) >>> r = nvmath.fft.fft(a, options=o) >>> r is a True
有关可用选项的完整列表,请参阅
FFTOptions
。默认情况下使用包的当前流,但可以显式地为 FFT 操作提供流。如果 FFT 操作数是在不同的流上计算的,则可以这样做,例如
>>> s = cp.cuda.Stream() >>> with s: ... a = cp.random.rand(*shape) + 1j * cp.random.rand(*shape) >>> r = nvmath.fft.fft(a, stream=s)
上述操作在流
s
上运行,并相对于输入计算进行排序。在 CPU 上创建一个 NumPy ndarray。
>>> import numpy as np >>> b = np.random.rand(*shape) + 1j * np.random.rand(*shape)
将 NumPy ndarray 提供给
fft()
,结果也是一个 NumPy ndarray>>> r = nvmath.fft.fft(b)
注解
此函数仅接受复数操作数进行 C2C 变换。如果用户希望对实数输入执行完整 FFT 变换,请将输入转换为相应的复数数据类型。
此函数是
FFT
的便捷包装器,专门用于单次使用。可以使用状态 API 通过FFT.execute()
中的默认direction
参数执行相同的计算。
更多示例可以在 nvmath/examples/fft 目录中找到。