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

tract_core/ops/
fft.rs
