fuel-transformers 0.19.1

// SPDX-License-Identifier: MIT OR Apache-2.0
//! Mimi streaming-capable 1-D convolution primitive (sub-port 1 of
//! port-mimi-conv.md).
//!
//! Ports the `StreamableConv1d` half of
//! `fuel_transformers::models::audio::mimi::conv` to the lazy-graph
//! API. The remaining variants
//! (`StreamableConvTranspose1d`, `ConvDownsample1d`,
//! `ConvTrUpsample1d`) compose on top of this primitive and ship in
//! the following sub-ports.
//!
//! # Differences from the eager API
//!
//! 1. **State as value, not interior mutability.** Eager carries the
//!    ring buffer in `&mut self`. The lazy port returns
//!    `(StreamConv1dState, Option<Tensor>)` from `step` so the
//!    streaming state composes naturally with the rest of the lazy
//!    encoder (no `&mut self` rippling through composition).
//! 2. **WeightNorm baked at load.** Eager's `conv1d_weight_norm`
//!    builds the effective weight once when the `Conv1d` is created.
//!    [`bake_weight_norm`] is the lazy equivalent — call it on the
//!    `(weight_g, weight_v)` pair you read out of safetensors before
//!    constructing the [`StreamableConv1dWeights`]. Once shipped,
//!    the runtime weight tensor is indistinguishable from a plain
//!    unnormalized conv weight.
//! 3. **TimeGroupNorm deferred.** No Mimi preset currently shipped to
//!    `fuel_transformers::models::audio::mimi::encodec` uses
//!    `Norm::TimeGroupNorm` in the convs; the variant is rejected at
//!    construction time when needed.
//!
//! # Padding modes
//!
//! - [`LazyPadMode::Constant`] zero-pad — the default for the SEANet
//!   encoder convs.
//! - [`LazyPadMode::Replicate`] edge-value pad — used by Mimi's
//!   `ConvDownsample1d` / `ConvTrUpsample1d`. Implemented as
//!   `narrow + repeat + concat` since `Op::Pad`'s Replicate mode
//!   isn't yet wired through the executor; that's an internal
//!   detail — callers see the same semantics.
//! - Reflect is rejected (matches eager's `bail!`).
//!
//! # Scope of sub-port 1
//!
//! - Dilation must equal 1 (the lazy `conv1d` primitive doesn't
//!   expose dilation; no Mimi preset that ships uses it).
//! - One-shot [`StreamableConv1dWeights::forward`] and chunked
//!   [`StreamableConv1dWeights::step`].
//! - WeightNorm via [`bake_weight_norm`]; Norm::Identity / None
//!   accepted.

use fuel_ir::Shape;
use fuel_ir::error::Result;
use fuel_tensor::lazy::Tensor;
use std::sync::Arc;

/// Padding mode for [`pad_last_1d`] / streaming-conv left context.
///
/// Mirrors eager `mimi::conv::PadMode` minus `Reflect`, which neither
/// eager Mimi nor the lazy executor currently support.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
/// ⚠️ KEEPS THE `Lazy` PREFIX ON PURPOSE — do not "finish" the 2026-08-19
/// prefix drop by renaming this. `PadMode` is ALREADY a live type in
/// `fuel-core/src/lazy_encodec.rs`, so here the prefix is NOT redundant: it is the only thing
/// distinguishing two types in one crate. CireSnave's ruling was to drop
/// `Lazy` *where it is redundant*, and this is one of exactly four names
/// (with LazyKvCache, LazyPadMode, LazyConv1dWeights,
/// LazyConvTranspose1dWeights) where that condition is not met.
pub enum LazyPadMode {
    /// Pad with zeros (eager `PadMode::Constant`).
    Constant,
    /// Repeat the edge value into the padded region
    /// (eager `PadMode::Replicate`, eager `pad_with_same`).
    Replicate,
}

/// Pad a `(B, C, T)` tensor along the temporal dim by `left` zeros /
/// edge-replications on the left and `right` on the right.
///
/// Replicate is implemented via `narrow(T, 0, 1).repeat(..., left)` +
/// concat — works against the existing `Op::Pad` (Constant only)
/// gap until the executor's Replicate path lands.
pub fn pad_last_1d(xs: &Tensor, left: usize, right: usize, mode: LazyPadMode) -> Result<Tensor> {
    if left == 0 && right == 0 {
        return Ok(xs.clone());
    }
    let dims = xs.shape();
    let dims = dims.dims();
    if dims.len() != 3 {
        return Err(fuel_ir::error::Error::Msg(format!(
            "pad_last_1d: input must be rank 3 (B, C, T), got rank {}",
            dims.len(),
        )));
    }
    match mode {
        LazyPadMode::Constant => Ok(xs.pad_with_zeros(2_usize, left, right)?),
        LazyPadMode::Replicate => {
            let (b, c, t) = (dims[0], dims[1], dims[2]);
            if t == 0 {
                return Err(fuel_ir::error::Error::Msg(
                    "pad_last_1d: cannot Replicate-pad an empty sequence".into(),
                ));
            }
            let mut out = xs.clone();
            let _ = (b, c);
            if left > 0 {
                let head = xs.narrow(2_usize, 0, 1)?;
                let left_block = head.repeat(Shape::from_dims(&[1, 1, left]))?;
                out = left_block.concat(&out, 2_usize)?;
            }
            if right > 0 {
                let tail = xs.narrow(2_usize, t - 1, 1)?;
                let right_block = tail.repeat(Shape::from_dims(&[1, 1, right]))?;
                out = out.concat(&right_block, 2_usize)?;
            }
            Ok(out)
        }
    }
}

/// Bake PyTorch-style weight normalization at load time.
///
/// PyTorch reparametrizes a conv weight as `W = g · v / ||v||`, with
/// `g` a per-output-channel scale of shape `(Cout, 1, 1)` and `v` a
/// raw weight of shape `(Cout, Cin/groups, K)`. The norm is computed
/// over the trailing two axes per output channel:
///
/// ```text
///     w[o, i, k] = g[o] · v[o, i, k] / sqrt(sum_{i,k} v[o, i, k]^2)
/// ```
///
/// Eager Mimi (`mimi::conv::conv1d_weight_norm`) bakes this once at
/// `Conv1d` construction time. The lazy port does the same here,
/// returning the effective weight in plain `(Cout, Cin/groups, K)`
/// row-major layout — indistinguishable from a non-WN-baked
/// checkpoint downstream.
// needless_range_loop here: the index feeds computed stride/offset addressing into a
// flat buffer (base = i*inner + j), which .iter() cannot express.
#[allow(clippy::needless_range_loop)]
pub fn bake_weight_norm(
    weight_g: &[f32],
    weight_v: &[f32],
    out_channels: usize,
    in_channels_per_group: usize,
    kernel_size: usize,
) -> Result<Vec<f32>> {
    if weight_g.len() != out_channels {
        return Err(fuel_ir::error::Error::Msg(format!(
            "bake_weight_norm: weight_g length {} != out_channels {}",
            weight_g.len(),
            out_channels,
        )));
    }
    let per_out = in_channels_per_group * kernel_size;
    if weight_v.len() != out_channels * per_out {
        return Err(fuel_ir::error::Error::Msg(format!(
            "bake_weight_norm: weight_v length {} != out_channels * \
             (in_channels_per_group * kernel_size) = {} * {} = {}",
            weight_v.len(),
            out_channels,
            per_out,
            out_channels * per_out,
        )));
    }
    let mut out = vec![0.0_f32; out_channels * per_out];
    for c in 0..out_channels {
        let base = c * per_out;
        let v_slice = &weight_v[base..base + per_out];
        let norm: f32 = v_slice.iter().map(|x| x * x).sum::<f32>().sqrt();
        if !norm.is_finite() || norm == 0.0 {
            return Err(fuel_ir::error::Error::Msg(format!(
                "bake_weight_norm: invalid norm {norm} at output channel {c}",
            )));
        }
        let scale = weight_g[c] / norm;
        for i in 0..per_out {
            out[base + i] = v_slice[i] * scale;
        }
    }
    Ok(out)
}

/// Weights + config for a streaming-capable 1-D conv.
///
/// `weight` is the *effective* kernel — call [`bake_weight_norm`]
/// before constructing this if the checkpoint stores the
/// `(weight_g, weight_v)` reparametrization pair.
#[derive(Debug, Clone)]
pub struct StreamableConv1dWeights {
    pub weight: Arc<[f32]>,
    pub bias: Option<Arc<[f32]>>,
    pub in_channels: usize,
    pub out_channels: usize,
    pub kernel_size: usize,
    pub stride: usize,
    pub groups: usize,
    pub causal: bool,
    pub pad_mode: LazyPadMode,
}

impl StreamableConv1dWeights {
    /// Validate shapes + store. `weight` must be
    /// `(out_channels, in_channels/groups, kernel_size)` row-major;
    /// `bias` (when `Some`) is `(out_channels,)`.
    #[allow(clippy::too_many_arguments)]
    pub fn new(
        weight: Arc<[f32]>,
        bias: Option<Arc<[f32]>>,
        in_channels: usize,
        out_channels: usize,
        kernel_size: usize,
        stride: usize,
        groups: usize,
        causal: bool,
        pad_mode: LazyPadMode,
    ) -> Result<Self> {
        if in_channels == 0 || out_channels == 0 || kernel_size == 0 || stride == 0 {
            return Err(fuel_ir::error::Error::Msg(format!(
                "StreamableConv1dWeights: in/out/kernel/stride must be > 0 \
                 (got {in_channels}, {out_channels}, {kernel_size}, {stride})",
            )));
        }
        if groups == 0 || !in_channels.is_multiple_of(groups) {
            return Err(fuel_ir::error::Error::Msg(format!(
                "StreamableConv1dWeights: in_channels {in_channels} must be divisible by groups {groups}",
            )));
        }
        if kernel_size < stride {
            return Err(fuel_ir::error::Error::Msg(format!(
                "StreamableConv1dWeights: kernel_size {kernel_size} must be >= stride {stride} \
                 (eager mimi::conv invariant)",
            )));
        }
        let expected_w = out_channels * (in_channels / groups) * kernel_size;
        if weight.len() != expected_w {
            return Err(fuel_ir::error::Error::Msg(format!(
                "StreamableConv1dWeights: weight length {} != \
                 out_channels * (in_channels/groups) * kernel_size = {expected_w}",
                weight.len(),
            )));
        }
        if let Some(b) = &bias
            && b.len() != out_channels
        {
            return Err(fuel_ir::error::Error::Msg(format!(
                "StreamableConv1dWeights: bias length {} != out_channels {}",
                b.len(),
                out_channels,
            )));
        }
        Ok(Self {
            weight,
            bias,
            in_channels,
            out_channels,
            kernel_size,
            stride,
            groups,
            causal,
            pad_mode,
        })
    }

    /// Effective kernel (= kernel for dilation=1 case).
    fn effective_kernel(&self) -> usize {
        self.kernel_size
    }

    /// Total padding the eager `StreamableConv1d::forward` applies
    /// (split causal vs symmetric below).
    fn padding_total(&self) -> usize {
        self.effective_kernel() - self.stride
    }

    /// Eager `get_extra_padding_for_conv1d`: extra right padding so
    /// the output length aligns to an integer number of strided
    /// frames. Returns 0 if input length already covers an exact
    /// number of frames.
    fn extra_padding(&self, t_in: usize) -> usize {
        let k = self.effective_kernel();
        let stride = self.stride;
        let total = self.padding_total();
        let n_frames = ((t_in + total).saturating_sub(k)) as f64 / stride as f64 + 1.0;
        let ideal = ((n_frames.ceil() as usize - 1) * stride + k).saturating_sub(total);
        ideal.saturating_sub(t_in)
    }

    fn build_weight_tensor(&self, anchor: &Tensor) -> fuel_ir::error::Result<Tensor> {
        // ⚠️ GAP-003 CATEGORY 3 — the ruled case, and the only kind that takes
        // `Result`. The DATA is `Arc::clone(&self.weight)`, loaded from a
        // CHECKPOINT. The SHAPE comes from CONFIG fields. There is no local
        // proof and there cannot be one: a mismatch means THE CHECKPOINT
        // DISAGREES WITH ITS CONFIG, which is not a bug in any caller's code.
        // It is invalid input arriving from outside the process, and a library
        // hands that back rather than aborting.
        //
        // The message names the CONFIG-IMPLIED count and the DIMS it came from,
        // because this text is the only thing a user debugging a bad checkpoint
        // will ever see. Modelled on `lazy_metavoice`'s speaker-embed guard,
        // which already does this correctly and was written without being asked.
        anchor
            .const_f32_like(
                Arc::clone(&self.weight),
                Shape::from_dims(&[
                    self.out_channels,
                    self.in_channels / self.groups,
                    self.kernel_size,
                ]),
            )
            .map_err(|e| {
                fuel_ir::error::Error::Msg(format!(
                    "conv weight from the checkpoint does not match the config: {e} \n                     (config implies out_channels * (in_channels / groups) * kernel_size)"
                ))
            })
    }

    fn build_bias_tensor(&self, anchor: &Tensor) -> fuel_ir::error::Result<Option<Tensor>> {
        // GAP-003 CATEGORY 3, same as the weight above: loaded bias data
        // against a config-derived length.
        self.bias
            .as_ref()
            .map(|b| {
                anchor
                    .const_f32_like(Arc::clone(b), Shape::from_dims(&[self.out_channels]))
                    .map_err(|e| {
                        fuel_ir::error::Error::Msg(format!(
                            "conv bias from the checkpoint does not match the config: {e} \n                             (config implies out_channels)"
                        ))
                    })
            })
            .transpose()
    }

    /// Run the streaming conv in one-shot mode. Matches eager
    /// `StreamableConv1d::forward` semantics: apply left/right
    /// padding according to causal/symmetric and pad-mode, then
    /// convolve at the underlying [`Tensor::conv1d`].
    ///
    /// Input shape `(B, in_channels, T)`. Output shape
    /// `(B, out_channels, T_out)` where
    /// `T_out = (T + pad_total + extra - kernel) / stride + 1`.
    pub fn forward(&self, xs: &Tensor) -> Result<Tensor> {
        let dims = xs.shape();
        let dims = dims.dims();
        if dims.len() != 3 || dims[1] != self.in_channels {
            return Err(fuel_ir::error::Error::Msg(format!(
                "StreamableConv1dWeights::forward: expected input (B, {}, T), got {dims:?}",
                self.in_channels,
            )));
        }
        let t_in = dims[2];
        let pad_total = self.padding_total();
        let extra = self.extra_padding(t_in);
        let padded = if self.causal {
            pad_last_1d(xs, pad_total, extra, self.pad_mode)?
        } else {
            let right = pad_total / 2;
            let left = pad_total - right;
            pad_last_1d(xs, left, right + extra, self.pad_mode)?
        };
        let w = self.build_weight_tensor(xs)?;
        let b = self.build_bias_tensor(xs)?;
        padded.conv1d(&w, b.as_ref(), self.stride, 0, self.groups)
    }
}

/// Persistent state for chunk-wise streaming inference through a
/// [`StreamableConv1dWeights`].
///
/// `left_pad_applied` mirrors eager `StreamableConv1d::left_pad_applied`:
/// the first chunk pays the initial left-context pad once; subsequent
/// chunks rely on `buf` to carry over the right context that wasn't
/// consumed by an integer number of strided frames in the previous
/// step.
///
/// `buf` is `Some(Tensor)` of shape `(B, in_channels, L)` once
/// any context has accumulated, otherwise `None`.
#[derive(Debug, Clone, Default)]
pub struct StreamConv1dState {
    pub buf: Option<Tensor>,
    pub left_pad_applied: bool,
}

impl StreamConv1dState {
    pub fn empty() -> Self {
        Self::default()
    }
}

impl StreamableConv1dWeights {
    /// Stream `xs` (one chunk of the input) through the conv and
    /// return the new state plus any output the chunk produces.
    ///
    /// `xs` shape: `(B, in_channels, L_chunk)`. The returned tensor
    /// (when `Some`) has shape `(B, out_channels, L_out)` where
    /// `L_out` is the number of strided frames the cumulative input
    /// has now produced past the previous step's end. When the
    /// chunk is smaller than the kernel context the call returns
    /// `None` and just grows the internal buffer.
    ///
    /// Calling `step` from the empty state on a length-T input
    /// followed by concatenating all returned chunks recovers
    /// [`Self::forward`] bit-for-bit up to floating-point
    /// reassociation in the conv kernel.
    pub fn step(
        &self,
        mut state: StreamConv1dState,
        xs: &Tensor,
    ) -> Result<(StreamConv1dState, Option<Tensor>)> {
        let dims = xs.shape();
        let dims = dims.dims();
        if dims.len() != 3 || dims[1] != self.in_channels {
            return Err(fuel_ir::error::Error::Msg(format!(
                "StreamableConv1dWeights::step: expected (B, {}, L), got {dims:?}",
                self.in_channels,
            )));
        }
        // On the first non-empty chunk pay the causal left-pad once.
        let xs_padded = if state.left_pad_applied {
            xs.clone()
        } else {
            state.left_pad_applied = true;
            let pad_total = self.padding_total();
            pad_last_1d(xs, pad_total, 0, self.pad_mode)?
        };
        // Combine with carried-over context.
        let combined = match state.buf.take() {
            None => xs_padded,
            Some(prev) => prev.concat(&xs_padded, 2_usize)?,
        };
        let combined_len = combined.shape().dims()[2];
        let kernel = self.effective_kernel();
        let stride = self.stride;
        let num_frames = (combined_len + stride).saturating_sub(kernel) / stride;
        if num_frames == 0 {
            state.buf = Some(combined);
            return Ok((state, None));
        }
        let offset = num_frames * stride;
        let carry = combined.narrow(2_usize, offset, combined_len - offset)?;
        let in_l = (num_frames - 1) * stride + kernel;
        let xs_in = combined.narrow(2_usize, 0, in_l)?;
        let w = self.build_weight_tensor(&xs_in)?;
        let b = self.build_bias_tensor(&xs_in)?;
        let y = xs_in.conv1d(&w, b.as_ref(), stride, 0, self.groups)?;
        state.buf = Some(carry);
        Ok((state, Some(y)))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use fuel_tensor::Device;

    fn const_xs(b: usize, c: usize, t: usize, src: &[f32]) -> Tensor {
        assert_eq!(src.len(), b * c * t);
        Tensor::from_f32(
            Arc::from(src.to_vec()),
            Shape::from_dims(&[b, c, t]),
            &Device::cpu(),
        )
        .unwrap()
    }

    #[test]
    fn bake_weight_norm_hand_computed() {
        // out=2, in=1, k=2 → v shape (2, 2) flat, g shape (2,).
        let g = vec![1.0_f32, 2.0];
        let v = vec![3.0_f32, 4.0, 6.0, 8.0];
        // norm[0] = sqrt(9 + 16) = 5 → w[0] = [3/5, 4/5] = [0.6, 0.8].
        // norm[1] = sqrt(36 + 64) = 10 → w[1] = [2 * 6/10, 2 * 8/10] = [1.2, 1.6].
        let w = bake_weight_norm(&g, &v, 2, 1, 2).unwrap();
        assert!((w[0] - 0.6).abs() < 1e-6);
        assert!((w[1] - 0.8).abs() < 1e-6);
        assert!((w[2] - 1.2).abs() < 1e-6);
        assert!((w[3] - 1.6).abs() < 1e-6);
    }

    #[test]
    fn bake_weight_norm_rejects_size_mismatch() {
        assert!(bake_weight_norm(&[1.0, 2.0], &[1.0, 2.0, 3.0], 2, 1, 2).is_err());
    }

    #[test]
    fn pad_last_1d_constant_matches_pad_with_zeros() {
        let xs = const_xs(1, 2, 3, &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]);
        let y = pad_last_1d(&xs, 1, 2, LazyPadMode::Constant).unwrap();
        let got = y.realize_f32();
        // Channel 0: [0, 1, 2, 3, 0, 0]; channel 1: [0, 4, 5, 6, 0, 0].
        assert_eq!(
            got,
            vec![0.0, 1.0, 2.0, 3.0, 0.0, 0.0, 0.0, 4.0, 5.0, 6.0, 0.0, 0.0]
        );
    }

    #[test]
    fn pad_last_1d_replicate_repeats_edges() {
        let xs = const_xs(1, 1, 3, &[1.0, 2.0, 3.0]);
        let y = pad_last_1d(&xs, 2, 1, LazyPadMode::Replicate).unwrap();
        let got = y.realize_f32();
        assert_eq!(got, vec![1.0, 1.0, 1.0, 2.0, 3.0, 3.0]);
    }

    #[test]
    fn pad_last_1d_replicate_multi_channel_preserves_channels() {
        // 2 channels: c0 = [10, 20, 30], c1 = [40, 50, 60].
        let xs = const_xs(1, 2, 3, &[10.0, 20.0, 30.0, 40.0, 50.0, 60.0]);
        let y = pad_last_1d(&xs, 1, 2, LazyPadMode::Replicate).unwrap();
        assert_eq!(y.shape().dims(), &[1, 2, 6]);
        let got = y.realize_f32();
        // c0: [10, 10, 20, 30, 30, 30]; c1: [40, 40, 50, 60, 60, 60].
        assert_eq!(
            got,
            vec![
                10.0, 10.0, 20.0, 30.0, 30.0, 30.0, 40.0, 40.0, 50.0, 60.0, 60.0, 60.0
            ],
        );
    }

    #[test]
    fn forward_kernel_1_stride_1_no_pad_is_identity_matmul() {
        // kernel=1, stride=1, pad_total = 0. With a small explicit
        // weight, the conv is just per-pixel matmul. Lets us
        // hand-check the forward shape + numerics.
        let weight: Arc<[f32]> = Arc::from(vec![2.0_f32, -1.0]); // (out=1, in=2, k=1)
        let cv =
            StreamableConv1dWeights::new(weight, None, 2, 1, 1, 1, 1, true, LazyPadMode::Constant)
                .unwrap();
        let xs = const_xs(1, 2, 3, &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]);
        let y = cv.forward(&xs).unwrap();
        assert_eq!(y.shape().dims(), &[1, 1, 3]);
        // y[t] = 2 * xs[ch=0,t] + (-1) * xs[ch=1,t]
        //      = [2*1 - 4, 2*2 - 5, 2*3 - 6] = [-2, -1, 0]
        let got = y.realize_f32();
        assert_eq!(got, vec![-2.0, -1.0, 0.0]);
    }

    fn stream_concat(
        cv: &StreamableConv1dWeights,
        xs: &Tensor,
        chunk_size: usize,
        t_total: usize,
    ) -> Tensor {
        let mut state = StreamConv1dState::empty();
        let mut pieces: Vec<Tensor> = Vec::new();
        let mut t = 0;
        while t < t_total {
            let len = chunk_size.min(t_total - t);
            let chunk = xs.narrow(2_usize, t, len).unwrap();
            let (new_state, y) = cv.step(state, &chunk).unwrap();
            state = new_state;
            if let Some(y) = y {
                pieces.push(y);
            }
            t += len;
        }
        let mut out = pieces[0].clone();
        for p in pieces.iter().skip(1) {
            out = out.concat(p, 2_usize).unwrap();
        }
        out
    }

    fn assert_close(a: &[f32], b: &[f32], tol: f32, label: &str) {
        assert_eq!(
            a.len(),
            b.len(),
            "{label}: length mismatch {} vs {}",
            a.len(),
            b.len()
        );
        for (i, (x, y)) in a.iter().zip(b.iter()).enumerate() {
            assert!(
                (x - y).abs() < tol,
                "{label}: idx {i}: streamed={x} one_shot={y} diff={}",
                (x - y).abs(),
            );
        }
    }

    #[test]
    fn streaming_chunk_by_1_equals_one_shot_kernel_3_stride_1_causal() {
        // Kernel=3, stride=1, in=2, out=2, causal. T=8.
        let w: Vec<f32> = (0..2 * 2 * 3).map(|i| (i as f32) * 0.1 - 0.5).collect();
        let cv = StreamableConv1dWeights::new(
            Arc::from(w),
            None,
            2,
            2,
            3,
            1,
            1,
            true,
            LazyPadMode::Constant,
        )
        .unwrap();
        let x_data: Vec<f32> = (0..2 * 8).map(|i| (i as f32) * 0.05).collect();
        let xs = const_xs(1, 2, 8, &x_data);
        let one_shot = cv.forward(&xs).unwrap().realize_f32();
        let streamed = stream_concat(&cv, &xs, 1, 8).realize_f32();
        assert_close(&streamed, &one_shot, 1e-5, "chunk_by_1");
    }

    #[test]
    fn streaming_chunk_by_3_equals_one_shot_kernel_3_stride_2_causal() {
        let w: Vec<f32> = (0..2 * 2 * 3).map(|i| 0.01 + (i as f32) * 0.07).collect();
        let cv = StreamableConv1dWeights::new(
            Arc::from(w),
            None,
            2,
            2,
            3,
            2,
            1,
            true,
            LazyPadMode::Constant,
        )
        .unwrap();
        let x_data: Vec<f32> = (0..2 * 9).map(|i| (i as f32) * 0.03 - 0.2).collect();
        let xs = const_xs(1, 2, 9, &x_data);
        let one_shot = cv.forward(&xs).unwrap().realize_f32();
        let streamed = stream_concat(&cv, &xs, 3, 9).realize_f32();
        // Streaming may produce one fewer output element than one-shot
        // because the extra right-padding eager applies in `forward`
        // isn't reproducible in pure streaming (no future input known).
        // Compare the common prefix that both pipelines computed.
        let n = streamed.len().min(one_shot.len());
        let per_ch = n / 2;
        // Channel-major layout: assert the first `per_ch` of each
        // channel agree.
        let cout = 2usize;
        let t_one = one_shot.len() / cout;
        let t_str = streamed.len() / cout;
        for c in 0..cout {
            for k in 0..per_ch {
                let a = streamed[c * t_str + k];
                let b = one_shot[c * t_one + k];
                assert!(
                    (a - b).abs() < 1e-5,
                    "stride2 ch={c} t={k}: streamed={a} one_shot={b}",
                );
            }
        }
    }

    #[test]
    fn rejects_kernel_smaller_than_stride() {
        let r = StreamableConv1dWeights::new(
            Arc::from(vec![0.0_f32; 2]),
            None,
            1,
            1,
            1,
            2,
            1,
            true,
            LazyPadMode::Constant,
        );
        assert!(r.is_err());
    }

    #[test]
    fn rejects_weight_size_mismatch() {
        let r = StreamableConv1dWeights::new(
            Arc::from(vec![0.0_f32; 3]),
            None,
            2,
            1,
            2,
            1,
            1,
            true,
            LazyPadMode::Constant,
        );
        assert!(r.is_err());
    }
}