tract_core/ops/
fft.rs

1use crate::internal::*;
2use num_complex::Complex;
3use rustfft::num_traits::{Float, FromPrimitive};
4use rustfft::{FftDirection, FftNum};
5use tract_data::itertools::Itertools;
6use tract_ndarray::Axis as NdAxis;
7
8#[derive(Clone, Debug, Hash, PartialEq, Eq)]
9pub struct Fft {
10    pub axis: usize,
11    pub inverse: bool,
12}
13
14impl Fft {
15    fn eval_t<T: Datum + FftNum + FromPrimitive + Float>(
16        &self,
17        tensor: &mut Tensor,
18    ) -> TractResult<()> {
19        let mut iterator_shape: TVec<usize> = tensor.shape().into();
20        iterator_shape.pop(); // last dim is [re, im]
21        iterator_shape[self.axis] = 1;
22        let len = tensor.shape()[self.axis];
23        let direction = if self.inverse { FftDirection::Inverse } else { FftDirection::Forward };
24        let fft = rustfft::FftPlanner::new().plan_fft(len, direction);
25        let mut tensor_plain = tensor.try_as_plain_mut()?;
26        let mut array = tensor_plain.to_array_view_mut::<T>()?;
27        let mut v = Vec::with_capacity(len);
28        for coords in tract_ndarray::indices(&*iterator_shape) {
29            v.clear();
30            let mut slice = array.slice_each_axis_mut(|ax| {
31                if ax.axis.index() == self.axis || ax.stride == 1 {
32                    // ax.stride == 1 => last dim
33                    (..).into()
34                } else {
35                    let c = coords[ax.axis.index()] as isize;
36                    (c..=c).into()
37                }
38            });
39            v.extend(slice.iter().tuples().map(|(r, i)| Complex::new(*r, *i)));
40            fft.process(&mut v);
41            slice
42                .iter_mut()
43                .zip(v.iter().flat_map(|cmpl| [cmpl.re, cmpl.im].into_iter()))
44                .for_each(|(s, v)| *s = v);
45        }
46        Ok(())
47    }
48}
49
50impl Op for Fft {
51    fn name(&self) -> StaticName {
52        "Fft".into()
53    }
54
55    fn info(&self) -> TractResult<Vec<String>> {
56        Ok(vec![if self.inverse { "inverse" } else { "forward" }.into()])
57    }
58
59    op_as_typed_op!();
60}
61
62impl EvalOp for Fft {
63    op_out_of_plan!();
64
65    fn eval(&self, _ctx: &EvalContext, inputs: TVec<TValue>) -> TractResult<TVec<TValue>> {
66        let mut tensor = args_1!(inputs).into_tensor();
67        match tensor.datum_type() {
68            DatumType::F16 => {
69                let mut temp = tensor.cast_to::<f32>()?.into_owned();
70                self.eval_t::<f32>(&mut temp)?;
71                tensor = temp.cast_to::<f16>()?.into_owned();
72            }
73            DatumType::F32 => self.eval_t::<f32>(&mut tensor)?,
74            DatumType::F64 => self.eval_t::<f64>(&mut tensor)?,
75            _ => bail!("FFT not implemented for type {:?}", tensor.datum_type()),
76        }
77        Ok(tvec!(tensor.into_tvalue()))
78    }
79}
80
81impl TypedOp for Fft {
82    fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
83        anyhow::ensure!(
84            inputs[0].rank() >= 2,
85            "Expect rank 2 (one for fft dimension, one for complex dimension"
86        );
87        anyhow::ensure!(
88            inputs[0].shape.last().unwrap() == &2.to_dim(),
89            "Fft operators expect inner (last) dimension to be 2 for real and imaginary part"
90        );
91        Ok(tvec!(inputs[0].without_value()))
92    }
93
94    fn axes_mapping(
95        &self,
96        inputs: &[&TypedFact],
97        _outputs: &[&TypedFact],
98    ) -> TractResult<AxesMapping> {
99        // Fft is rank-preserving but it is NOT axes-natural: two axes do
100        // not map 1-to-1 from input to output and must be declared as a
101        // separate input-only and output-only axis.
102        //
103        //   - the FFT axis (`self.axis`): every output sample along it
104        //     depends on every input sample, so the axis cannot be
105        //     sliced or streamed.
106        //   - the trailing complex axis (`rank - 1`): the FFT mixes the
107        //     real and imaginary parts, so re/im do not map 1-to-1.
108        //
109        // Splitting them is exactly what makes the generic pulse fallback
110        // bail when asked to track a streaming axis through the FFT or
111        // complex axis, while every genuine batch axis stays 1-to-1 and
112        // is handled by the per-pulse `PulseWrappingOp`. No dedicated
113        // `Fft` pulsifier is needed.
114        let rank = inputs[0].rank();
115        let complex_axis = rank - 1;
116        let mut axes = tvec!();
117        let mut alphabet = 'a'..;
118        for i in 0..rank {
119            if i == self.axis || i == complex_axis {
120                axes.push(crate::axes::Axis::new(alphabet.next().unwrap(), 1, 1).input(0, i));
121                axes.push(crate::axes::Axis::new(alphabet.next().unwrap(), 1, 1).output(0, i));
122            } else {
123                axes.push(
124                    crate::axes::Axis::new(alphabet.next().unwrap(), 1, 1).input(0, i).output(0, i),
125                );
126            }
127        }
128        AxesMapping::new(1, 1, axes)
129    }
130
131    as_op!();
132}
133
134#[derive(Clone, Debug, Hash, PartialEq, Eq)]
135pub struct Stft {
136    pub axis: usize,
137    pub frame: usize,
138    pub stride: usize,
139    pub window: Option<Arc<Tensor>>,
140}
141
142impl Stft {
143    fn eval_t<T: Datum + FftNum + FromPrimitive + Float>(
144        &self,
145        input: &Tensor,
146    ) -> TractResult<Tensor> {
147        let mut iterator_shape: TVec<usize> = input.shape().into();
148        iterator_shape.pop(); // [re,im]
149        iterator_shape[self.axis] = 1;
150        let mut output_shape: TVec<usize> = input.shape().into();
151        let frames = (input.shape()[self.axis] - self.frame) / self.stride + 1;
152        output_shape.insert(self.axis, frames);
153        output_shape[self.axis + 1] = self.frame;
154        let mut output = unsafe { Tensor::uninitialized::<T>(&output_shape)? };
155        let fft = rustfft::FftPlanner::new().plan_fft_forward(self.frame);
156        let input = input.to_plain_array_view::<T>()?;
157        let mut output_plain = output.try_as_plain_mut()?;
158        let mut oview = output_plain.to_array_view_mut::<T>()?;
159        let mut v = Vec::with_capacity(self.frame);
160        for coords in tract_ndarray::indices(&*iterator_shape) {
161            let islice = input.slice_each_axis(|ax| {
162                if ax.axis.index() == self.axis || ax.stride == 1 {
163                    (..).into()
164                } else {
165                    let c = coords[ax.axis.index()] as isize;
166                    (c..=c).into()
167                }
168            });
169            let mut oslice = oview.slice_each_axis_mut(|ax| {
170                if ax.stride == 1 {
171                    (..).into()
172                } else if ax.axis.index() < self.axis {
173                    let c = coords[ax.axis.index()] as isize;
174                    (c..=c).into()
175                } else if ax.axis.index() == self.axis || ax.axis.index() == self.axis + 1 {
176                    (..).into()
177                } else {
178                    let c = coords[ax.axis.index() - 1] as isize;
179                    (c..=c).into()
180                }
181            });
182            let signal: Vec<Complex<T>> =
183                islice.iter().tuples().map(|(re, im)| Complex::new(*re, *im)).collect();
184            for f in 0..frames {
185                v.clear();
186                v.extend_from_slice(&signal[self.stride * f..self.stride * f + self.frame]);
187                if let Some(win) = &self.window {
188                    let win = win.try_as_plain()?.as_slice::<T>()?;
189                    // symmetric padding in case window is smaller than frames (aka n fft)
190                    let pad_left = (self.frame - win.len()) / 2;
191                    v.iter_mut().enumerate().for_each(|(ix, v)| {
192                        *v = if ix < pad_left || ix >= pad_left + win.len() {
193                            Complex::new(T::zero(), T::zero())
194                        } else {
195                            *v * Complex::new(win[ix - pad_left], T::zero())
196                        }
197                    });
198                }
199                fft.process(&mut v);
200                oslice
201                    .index_axis_mut(NdAxis(self.axis), f)
202                    .iter_mut()
203                    .zip(v.iter().flat_map(|cmpl| [cmpl.re, cmpl.im].into_iter()))
204                    .for_each(|(s, v)| *s = v);
205            }
206        }
207        Ok(output)
208    }
209}
210
211impl Op for Stft {
212    fn name(&self) -> StaticName {
213        "STFT".into()
214    }
215
216    op_as_typed_op!();
217}
218
219impl EvalOp for Stft {
220    op_out_of_plan!();
221
222    fn eval(&self, _ctx: &EvalContext, inputs: TVec<TValue>) -> TractResult<TVec<TValue>> {
223        let input = args_1!(inputs);
224        let output = match input.datum_type() {
225            DatumType::F16 => {
226                let temp = input.cast_to::<f32>()?;
227                self.eval_t::<f32>(&temp)?.cast_to::<f16>()?.into_owned()
228            }
229            DatumType::F32 => self.eval_t::<f32>(&input)?,
230            DatumType::F64 => self.eval_t::<f64>(&input)?,
231            _ => bail!("FFT not implemented for type {:?}", input.datum_type()),
232        };
233        Ok(tvec!(output.into_tvalue()))
234    }
235}
236
237impl TypedOp for Stft {
238    fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
239        anyhow::ensure!(
240            inputs[0].rank() >= 2,
241            "Expect rank 2 (one for fft dimension, one for complex dimension"
242        );
243        anyhow::ensure!(
244            inputs[0].shape.last().unwrap() == &2.to_dim(),
245            "Fft operators expect inner (last) dimension to be 2 for real and imaginary part"
246        );
247        let mut shape = inputs[0].shape.to_tvec();
248        let frames = (inputs[0].shape[self.axis].clone() - self.frame) / self.stride + 1;
249        shape[self.axis] = frames;
250        shape.insert(self.axis + 1, self.frame.to_dim());
251        Ok(tvec!(inputs[0].datum_type.fact(shape)))
252    }
253
254    fn axes_mapping(
255        &self,
256        inputs: &[&TypedFact],
257        _outputs: &[&TypedFact],
258    ) -> TractResult<crate::axes::AxesMapping> {
259        // Stft is NOT rank-preserving: it inserts a frame axis at
260        // `axis + 1`. The mapping is:
261        //   - axes 0..self.axis (leading dims): 1-to-1 input <-> output.
262        //   - input axis `self.axis` (the time axis) <-> output axis
263        //     `self.axis` (now the n_frames axis -- same position, the
264        //     dim shrinks from `T` to `(T - frame) / stride + 1`).
265        //   - output axis `self.axis + 1` (the inserted frame axis):
266        //     output-only, no input correspondence.
267        //   - input axes `self.axis + 1..rank` (trailing dims incl.
268        //     the complex pair) <-> output axes `self.axis + 2..rank+1`
269        //     (shifted right by 1 to make room for the frame axis).
270        //
271        // Without this mapping the generic `PulseWrappingOp` fallback
272        // bails with "could not track pulsing axis" the moment a user
273        // streams a non-time axis through STFT (typical pattern: a
274        // batched STFT pipeline that streams the batch axis).
275        let in_rank = inputs[0].rank();
276        let mut axes = tvec!();
277        let mut alphabet = 'a'..;
278        for i in 0..in_rank {
279            let out_axis = if i <= self.axis { i } else { i + 1 };
280            axes.push(
281                crate::axes::Axis::new(alphabet.next().unwrap(), 1, 1)
282                    .input(0, i)
283                    .output(0, out_axis),
284            );
285        }
286        // Inserted frame axis (output-only).
287        axes.push(crate::axes::Axis::new(alphabet.next().unwrap(), 1, 1).output(0, self.axis + 1));
288        crate::axes::AxesMapping::new(1, 1, axes)
289    }
290
291    as_op!();
292}
tract_core/ops/fft.rs

tract_core/ops/
fft.rs
