rlx_ir/ops/
audio_ops.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3// SPDX-License-Identifier: MIT OR Apache-2.0
4
5//! Audio frontend graph builders.
6
7use crate::{Graph, NodeId, Op};
8
9impl Graph {
10    /// Log-mel spectrogram (Whisper-style) from RLX FFT block-layout spectrum.
11    ///
12    /// * **spectrum** — `[..., 2*n_fft]` with re plane then im plane (same as `Op::Fft` output).
13    /// * **filters** — `[n_mels, n_bins]` mel filterbank (`n_bins = n_fft/2 + 1`).
14    ///
15    /// Output: `[..., n_mels]`.
16    pub fn log_mel(&mut self, spectrum: NodeId, filters: NodeId) -> NodeId {
17        let spec_shape = self.shape(spectrum).clone();
18        let filt_shape = self.shape(filters).clone();
19        let out = crate::audio::log_mel_output_shape(&spec_shape, &filt_shape)
20            .unwrap_or_else(|e| panic!("log_mel shape error: {e}"));
21        self.push(Op::LogMel, vec![spectrum, filters], out, None)
22    }
23
24    /// VJP of [`log_mel`] w.r.t. `spectrum`.
25    pub fn log_mel_backward(&mut self, spectrum: NodeId, filters: NodeId, dy: NodeId) -> NodeId {
26        let spec_shape = self.shape(spectrum).clone();
27        self.push(
28            Op::LogMelBackward,
29            vec![spectrum, filters, dy],
30            spec_shape,
31            None,
32        )
33    }
34
35    /// Top-K Welch peaks from RLX FFT block-layout segment spectra.
36    ///
37    /// * **spectrum** — `[batch * n_segments, 2*n_fft]`
38    /// * **k** — spikes per batch row
39    /// * **n_segments** — Welch segments averaged per row
40    ///
41    /// Output: `[batch, k*2]` interleaved `(bin, power)`.
42    pub fn welch_peaks(&mut self, spectrum: NodeId, k: usize, n_segments: usize) -> NodeId {
43        let spec_shape = self.shape(spectrum).clone();
44        let out = crate::audio::welch_peaks_output_shape(&spec_shape, k, n_segments)
45            .unwrap_or_else(|e| panic!("welch_peaks shape error: {e}"));
46        self.push(Op::WelchPeaks { k, n_segments }, vec![spectrum], out, None)
47    }
48}
rlx_ir/ops/audio_ops.rs

rlx_ir/ops/
audio_ops.rs
