rlx-ir 0.2.3

// RLX — versatile ML compiler + runtime.
// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, version 3.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program. If not, see <https://www.gnu.org/licenses/>.

//! Audio frontend graph builders.

use crate::{Graph, NodeId, Op};

impl Graph {
    /// Log-mel spectrogram (Whisper-style) from RLX FFT block-layout spectrum.
    ///
    /// * **spectrum** — `[..., 2*n_fft]` with re plane then im plane (same as `Op::Fft` output).
    /// * **filters** — `[n_mels, n_bins]` mel filterbank (`n_bins = n_fft/2 + 1`).
    ///
    /// Output: `[..., n_mels]`.
    pub fn log_mel(&mut self, spectrum: NodeId, filters: NodeId) -> NodeId {
        let spec_shape = self.shape(spectrum).clone();
        let filt_shape = self.shape(filters).clone();
        let out = crate::audio::log_mel_output_shape(&spec_shape, &filt_shape)
            .unwrap_or_else(|e| panic!("log_mel shape error: {e}"));
        self.push(Op::LogMel, vec![spectrum, filters], out, None)
    }

    /// VJP of [`log_mel`] w.r.t. `spectrum`.
    pub fn log_mel_backward(&mut self, spectrum: NodeId, filters: NodeId, dy: NodeId) -> NodeId {
        let spec_shape = self.shape(spectrum).clone();
        self.push(
            Op::LogMelBackward,
            vec![spectrum, filters, dy],
            spec_shape,
            None,
        )
    }
}