tenferro_fft/
spec.rs

1use tenferro_tensor::DType;
2
3/// FFT normalization convention.
4///
5/// `Backward` matches NumPy, JAX, and PyTorch defaults: the forward transform
6/// is unscaled and the inverse transform is scaled by `1 / n`.
7///
8/// # Examples
9///
10/// ```
11/// use tenferro_fft::FftNorm;
12///
13/// assert_eq!(FftNorm::default(), FftNorm::Backward);
14/// ```
15#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
16pub enum FftNorm {
17    /// Scale inverse transforms by `1 / n`.
18    #[default]
19    Backward,
20    /// Scale forward transforms by `1 / n`.
21    Forward,
22    /// Scale both forward and inverse transforms by `1 / sqrt(n)`.
23    Ortho,
24}
25
26impl FftNorm {
27    #[cfg(feature = "autodiff")]
28    pub(crate) fn c2c_adjoint(self) -> Self {
29        match self {
30            Self::Backward => Self::Forward,
31            Self::Forward => Self::Backward,
32            Self::Ortho => Self::Ortho,
33        }
34    }
35}
36
37/// One-dimensional FFT operation requested from an [`FftBackend`](crate::FftBackend).
38///
39/// # Examples
40///
41/// ```
42/// use tenferro_fft::FftOperation;
43///
44/// assert_ne!(FftOperation::C2cForward, FftOperation::C2cInverse);
45/// assert_ne!(FftOperation::R2cFull, FftOperation::R2cOnesided);
46/// ```
47#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
48pub enum FftOperation {
49    /// Forward complex-to-complex FFT.
50    C2cForward,
51    /// Inverse complex-to-complex FFT.
52    C2cInverse,
53    /// Real-to-complex FFT returning the full complex spectrum.
54    R2cFull,
55    /// Real-to-complex FFT returning the non-negative-frequency spectrum.
56    R2cOnesided,
57    /// Complex-to-real FFT consuming a one-sided Hermitian spectrum.
58    C2r,
59}
60
61impl FftOperation {
62    pub(crate) const fn is_c2c(self) -> bool {
63        matches!(self, Self::C2cForward | Self::C2cInverse)
64    }
65
66    pub(crate) const fn is_forward(self) -> bool {
67        matches!(self, Self::C2cForward | Self::R2cFull | Self::R2cOnesided)
68    }
69
70    pub(crate) const fn is_onesided(self) -> bool {
71        matches!(self, Self::R2cOnesided)
72    }
73}
74
75/// Validated, backend-neutral description of one FFT request.
76///
77/// Backends receive this value only after the public FFT surface has validated
78/// dtype, rank, axis, requested length, and inverse-spectrum shape. The layout
79/// requirement is explicit so a backend can reject unsupported storage without
80/// silently materializing or transferring the input.
81///
82/// # Examples
83///
84/// ```
85/// use tenferro_fft::{FftOperation, FftPlanSpec};
86///
87/// fn inspect(spec: &FftPlanSpec) {
88///     let operation: FftOperation = spec.operation();
89///     assert_eq!(operation, spec.operation());
90///     assert!(spec.requires_compact_column_major());
91/// }
92/// ```
93#[derive(Clone, Debug, Eq, PartialEq)]
94pub struct FftPlanSpec {
95    operation: FftOperation,
96    normalized_axis: usize,
97    requested_len: Option<usize>,
98    norm: FftNorm,
99    input_dtype: DType,
100    input_shape: Vec<usize>,
101    requires_compact_column_major: bool,
102}
103
104impl FftPlanSpec {
105    pub(crate) fn new(
106        operation: FftOperation,
107        normalized_axis: usize,
108        requested_len: Option<usize>,
109        norm: FftNorm,
110        input_dtype: DType,
111        input_shape: Vec<usize>,
112    ) -> Self {
113        Self {
114            operation,
115            normalized_axis,
116            requested_len,
117            norm,
118            input_dtype,
119            input_shape,
120            requires_compact_column_major: true,
121        }
122    }
123
124    /// Return the transform operation.
125    pub const fn operation(&self) -> FftOperation {
126        self.operation
127    }
128
129    /// Return the non-negative, validated axis.
130    pub const fn normalized_axis(&self) -> usize {
131        self.normalized_axis
132    }
133
134    /// Return the caller-requested transform length, if one was supplied.
135    pub const fn requested_len(&self) -> Option<usize> {
136        self.requested_len
137    }
138
139    /// Return the requested normalization convention.
140    pub const fn norm(&self) -> FftNorm {
141        self.norm
142    }
143
144    /// Return the validated input dtype.
145    pub const fn input_dtype(&self) -> DType {
146        self.input_dtype
147    }
148
149    /// Return the validated input shape.
150    pub fn input_shape(&self) -> &[usize] {
151        &self.input_shape
152    }
153
154    /// Return whether execution requires compact column-major storage.
155    pub const fn requires_compact_column_major(&self) -> bool {
156        self.requires_compact_column_major
157    }
158}
tenferro_fft/spec.rs

tenferro_fft/
spec.rs
