nvmath-python 0.9.0__cp314-cp314-win_amd64.whl

nvmath/distributed/fft/__init__.py

@@ -0,0 +1,7 @@
+# Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from ..distribution import Slab as Slab
+from ._configuration import *  # noqa: F403
+from .fft import *  # noqa: F403

nvmath/distributed/fft/_configuration.py

@@ -0,0 +1,79 @@
+# Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+__all__ = ["FFTDirection", "FFTOptions"]
+
+from dataclasses import dataclass
+from enum import IntEnum
+from logging import Logger
+from typing import Literal
+
+
+@dataclass
+class FFTOptions:
+    """
+    A data class for providing options to the :class:`FFT` object and the family of wrapper
+    functions :func:`fft`, :func:`ifft`, :func:`rfft`, and :func:`irfft`.
+
+    Attributes:
+        fft_type: The type of FFT to perform, available options include ``'C2C'``,
+            ``'C2R'``, and ``'R2C'``. The default is ``'C2C'`` for complex input and
+            ``'R2C'`` for real input.
+
+        reshape: Reshape the output distribution to the same slab distribution used by the
+            input. This only applies when using a Slab distribution. The default is `True`.
+
+        last_axis_parity: For complex-to-real FFT (corresponding to ``fft_type='C2R'``),
+            specify whether the global size of the last axis in the result should be even
+            or odd. The even size is calculated as :math:`2 * (m - 1)`, where :math:`m` is
+            the size of the last axis of the operand, and the odd size is calculated as
+            :math:`2 * (m - 1) + 1`. The specified value should be either ``'even'`` or
+            ``'odd'``, with the default being ``'even'``.
+
+        logger (logging.Logger): Python Logger object. The root logger will be used if a
+            logger object is not provided.
+
+        blocking: A flag specifying the behavior of the execution functions and methods,
+            such as :func:`fft` and :meth:`FFT.execute`. When ``blocking`` is `True`, the
+            execution methods do not return until the operation is complete. When
+            ``blocking`` is ``"auto"``, the methods return immediately when the input tensor
+            is on the GPU. The execution methods always block when the input tensor is on
+            the CPU, to ensure that the user doesn't inadvertently use the result before it
+            becomes available. The default is ``"auto"``.
+
+    .. seealso::
+        :class:`FFT`, :func:`fft`, :func:`ifft`, :func:`rfft`, and :func:`irfft`.
+    """
+
+    fft_type: Literal["C2C", "C2R", "R2C"] | None = None
+    reshape: bool = True
+    last_axis_parity: Literal["even", "odd"] = "even"
+    logger: Logger | None = None
+    blocking: Literal[True, "auto"] = "auto"
+
+    def __post_init__(self):
+        valid_fft_types = [None, "C2C", "C2R", "R2C"]
+        if self.fft_type not in valid_fft_types:
+            raise ValueError(f"The value specified for 'fft_type' must be one of {valid_fft_types}.")
+
+        if not isinstance(self.reshape, bool):
+            raise ValueError("The value specified for 'reshape' must be of type bool (True or False).")
+
+        valid_last_axis_parity = ["even", "odd"]
+        if self.last_axis_parity not in valid_last_axis_parity:
+            raise ValueError(f"The value specified for 'last_axis_parity' must be one of {valid_last_axis_parity}.")
+
+        if self.blocking not in (True, "auto"):
+            raise ValueError("The value specified for 'blocking' must be either True or 'auto'.")
+
+
+class FFTDirection(IntEnum):
+    """An IntEnum class specifying the direction of the transform.
+
+    .. seealso::
+        :meth:`FFT.execute`, :func:`fft`
+    """
+
+    FORWARD = -1
+    INVERSE = 1