asry 0.1.0

//! ONNX encode + log-softmax stage of the alignment algorithm.
//!
//! Only [`encode_log_softmax`] itself needs `ort` (the ONNX
//! session call) and is gated on `feature = "alignment"`. Its
//! output type [`LogProbsTV`] and the log-softmax / validation
//! helpers below it are ort-free and reachable under the
//! `emissions` feature too — callers with their own acoustic
//! encoder (e.g. a CoreML wav2vec2 port) construct a `LogProbsTV`
//! from their own model output via [`LogProbsTV::new`] and feed it
//! straight into `asry::emissions::align_emissions`.

#[cfg(feature = "alignment")]
use ort::{
  session::{RunOptions, Session},
  value::{Shape, Tensor},
};
use smol_str::{SmolStr, format_smolstr};

use super::errors::{EmissionsError, EmissionsFailure};
use crate::types::{AlignmentError, AlignmentFailure, Lang, WorkFailure};

// NOTE on the (1, T) reshape `encode_log_softmax` uses below: the
// plan's literal pseudocode uses `ndarray::Array2::from_shape_vec((1,
// T), …)`, but wiring a direct `ndarray` dependency alongside ort's
// OWN internal `ndarray` re-export (`ort`'s `features = ["ndarray"]`
// flag) would pull in two independent copies of the crate —
// `Tensor::from_array(Array<T, D>)` only resolves for ort's own
// `ndarray` version, so a caller-side `ndarray::Array2` doesn't
// satisfy the trait bound. asry carried a direct `ndarray`
// dependency for exactly this reshape at one point; it turned out
// unused (this file never called into it) and was dropped. We use
// ort's version-agnostic `OwnedTensorArrayData for (D, Vec<T>)` impl
// instead (`Tensor::from_array((shape, v))`), which is exactly the
// constructor the ort docs use in their session-input examples. This
// keeps the (1, T) reshape semantically identical without forcing a
// cross-version `ndarray` bridge or an unused direct dependency.

/// The shape/indexing arm of [`LogProbsError`] — the
/// [`LogProbsError::Shape`] payload [`LogProbsTV::new`] returns when
/// it rejects the `(t, v, data.len())` triple: either `t * v !=
/// data.len()` (the flat buffer's length doesn't match the declared
/// `(T, V)` shape, including the overflow case where `t * v` doesn't
/// fit `usize`), or `v == 0` (a CTC vocabulary must contain at least
/// the blank token, so a zero-length vocab axis is never valid,
/// regardless of `t`). The value-domain arm — a value outside the
/// log-probability domain (finite ∧ `≤ 0`) — is
/// [`LogProbsValueError`]. Message text is
/// chosen by [`Display`](core::fmt::Display) rather than
/// `thiserror`'s `#[error(...)]` shorthand because the two shape
/// failures need different wording — reusing the shape-mismatch
/// wording for the zero-vocab case would claim `t * v != data.len()`
/// when it does not.
#[derive(Clone, Copy, Debug, PartialEq, Eq, thiserror::Error)]
pub struct LogProbsShapeError {
  t: usize,
  v: usize,
  data_len: usize,
}

impl core::fmt::Display for LogProbsShapeError {
  fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
    if self.v == 0 {
      write!(
        f,
        "LogProbsTV has a zero-length vocab dim: t={}, v=0, data.len()={} \
 (a CTC vocabulary must contain at least the blank token)",
        self.t, self.data_len
      )
    } else {
      write!(
        f,
        "LogProbsTV shape mismatch: t={}, v={}, data.len()={} (expected \
 data.len() == t * v)",
        self.t, self.v, self.data_len
      )
    }
  }
}

impl LogProbsShapeError {
  const fn new(t: usize, v: usize, data_len: usize) -> Self {
    Self { t, v, data_len }
  }

  /// Time dimension the caller supplied to [`LogProbsTV::new`].
  #[must_use]
  pub const fn t(&self) -> usize {
    self.t
  }

  /// Vocab dimension the caller supplied to [`LogProbsTV::new`].
  #[must_use]
  pub const fn v(&self) -> usize {
    self.v
  }

  /// Actual length of the flat buffer the caller supplied to
  /// [`LogProbsTV::new`].
  #[must_use]
  pub const fn data_len(&self) -> usize {
    self.data_len
  }
}

/// How a value fell outside the log-probability domain (finite ∧
/// `≤ 0`) that the value-domain scan in [`LogProbsTV::new`]
/// enforces. Reported by [`LogProbsValueError`] in place of the raw
/// `f32` because a raw `f32` is not `Eq` (`NaN != NaN`) — a class
/// keeps the error `Copy + Eq`, matching its [`LogProbsShapeError`]
/// sibling.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum LogProbsValueClass {
  /// `f32::NAN`.
  Nan,
  /// `f32::INFINITY` (`+∞`).
  PosInf,
  /// `f32::NEG_INFINITY` (`−∞`).
  NegInf,
  /// A finite value strictly greater than `0.0`. Log-probabilities
  /// are `≤ 0` (`log(p)` for `p ∈ (0, 1]`), so a positive value is
  /// not a log-probability: exponentiated by the DP it leaves
  /// `[0, 1]` (`f32::MAX.exp() = +∞`, `(1e-7).exp() ≈ 1.0000001`).
  Positive,
}

impl core::fmt::Display for LogProbsValueClass {
  fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
    f.write_str(match self {
      Self::Nan => "NaN",
      Self::PosInf => "+Inf",
      Self::NegInf => "-Inf",
      Self::Positive => "positive",
    })
  }
}

/// The value-domain arm of [`LogProbsError`] — the
/// [`LogProbsError::Value`] payload [`LogProbsTV::new`] returns when
/// `data` holds a value outside the log-probability domain (finite ∧
/// `≤ 0`). Locates the first offending element by `frame` (row) and
/// `vocab_index` (column) and records its [`LogProbsValueClass`]
/// rather than the raw `f32`, so the error stays `Copy + Eq` like
/// its [`LogProbsShapeError`] sibling. See [`LogProbsTV::new`] for
/// the full value-domain rule (why `NaN`, `±∞`, and a finite
/// positive value are all rejected).
///
/// The offending element need not be non-finite — a finite value
/// `> 0.0` is also out of domain, and is reported as
/// [`LogProbsValueClass::Positive`], not `NaN`/`PosInf`/`NegInf`.
/// Pinned by `new_rejects_tiny_positive_value`.
#[derive(Clone, Copy, Debug, PartialEq, Eq, thiserror::Error)]
#[error(
  "log-probability out of domain (finite and ≤ 0) at frame {frame}, vocab {vocab_index}: {class}"
)]
pub struct LogProbsValueError {
  frame: usize,
  vocab_index: usize,
  class: LogProbsValueClass,
}

impl LogProbsValueError {
  const fn new(frame: usize, vocab_index: usize, class: LogProbsValueClass) -> Self {
    Self {
      frame,
      vocab_index,
      class,
    }
  }

  /// Frame (row) index of the first out-of-domain element.
  #[must_use]
  pub const fn frame(&self) -> usize {
    self.frame
  }

  /// Vocab (column) index of the first out-of-domain element.
  #[must_use]
  pub const fn vocab_index(&self) -> usize {
    self.vocab_index
  }

  /// Which value-domain class the offending element fell into.
  #[must_use]
  pub const fn class(&self) -> LogProbsValueClass {
    self.class
  }
}

/// Everything [`LogProbsTV::new`] can reject, in one type so the
/// constructor's full input contract is enumerable at a glance.
/// [`Shape`](Self::Shape) carries the dimension/indexing rules
/// (product mismatch, `t * v` overflow, zero-length vocab axis);
/// [`Value`](Self::Value) carries the value-domain rule (a value
/// outside the log-probability domain: non-finite, or finite but
/// `> 0`). Both arms forward their `Display` to the wrapped error
/// via `#[error(transparent)]`.
#[derive(Clone, Copy, Debug, PartialEq, Eq, thiserror::Error)]
pub enum LogProbsError {
  /// The `(t, v, data.len())` triple is inconsistent: `t * v !=
  /// data.len()`, `t * v` overflows `usize`, or `v == 0`. See
  /// [`LogProbsShapeError`].
  #[error(transparent)]
  Shape(LogProbsShapeError),
  /// `data` contained a value outside the log-probability domain
  /// (non-finite, or finite but `> 0`). See [`LogProbsValueError`].
  #[error(transparent)]
  Value(LogProbsValueError),
}

/// Output of `encode_log_softmax` — or, under the `emissions`
/// feature, of the caller's own acoustic encoder. `pub` so both
/// the `feature = "bench-internals"` re-export and
/// `asry::emissions` can reach it.
pub struct LogProbsTV {
  /// Time dimension (number of wav2vec2 output frames).
  t: usize,
  /// Vocab dimension.
  v: usize,
  /// Flat row-major `(T, V)` log-probabilities. Index with
  /// `[t * v_dim + v_idx]`.
  data: Vec<f32>,
}

impl LogProbsTV {
  /// Construct from explicit dimensions + a flat row-major `(T, V)`
  /// buffer, validating the **complete** input contract below. This
  /// is the public, validating entry point an `emissions`-feature
  /// caller (its own acoustic encoder, no ort session) uses to hand
  /// its output to the alignment pipeline. The internal ort path
  /// builds a `LogProbsTV` by a different route — a struct literal
  /// after [`log_softmax_with_finite_guard`] — and never runs this
  /// constructor or its scan.
  ///
  /// The validated contract, checked in order:
  ///
  /// 1. **Vocab axis non-empty** — `v != 0` (checked first, for
  ///    every `t` including `t == 0`: a CTC vocabulary must contain
  ///    at least the blank token, so a zero-length vocab axis is
  ///    never valid).
  /// 2. **Shape / product / overflow** — `t.checked_mul(v) ==
  ///    Some(data.len())`. `checked_mul` rejects a `t`/`v` pair
  ///    whose product overflows `usize` rather than letting it wrap
  ///    to a small product that spuriously matches a small
  ///    `data.len()`.
  /// 3. **Value domain** — every element of `data` is a valid
  ///    log-probability: finite (`f32::is_finite`) **and** `≤ 0.0`.
  ///    `NaN`, `+∞`, `−∞`, and any finite value `> 0.0` are all
  ///    rejected. This is exactly the mathematical domain of a
  ///    log-probability (`log(p)` for `p ∈ (0, 1]` is finite and
  ///    `≤ 0`), and exactly the domain the internal ort path
  ///    produces: [`log_softmax_with_finite_guard`] emits
  ///    `lp = (x − max) − ln Σ exp(x − max)`; the `max` element
  ///    contributes `exp(0) = 1` to the sum, so `Σ ≥ 1` and
  ///    `ln Σ ≥ 0` fp-wise, hence every output is `(≤ 0) − (≥ 0) ≤
  ///    0.0` exactly. Enforcing the same `≤ 0` bound at this
  ///    external seam keeps its numeric domain identical to the one
  ///    the DP was built and tested against — no wider.
  ///
  ///    Why each rejection matters. The DP reads an emission as
  ///    `exp(lp)`. A `NaN` seeds a `NaN` word confidence
  ///    (`f32::max`-based trellis comparisons silently drop it, the
  ///    end-cell finiteness guard checks a *different* cell than the
  ///    backtrack seed reads); a finite `> 0` value exponentiates
  ///    *out of* `[0, 1]` (`f32::MAX.exp() = +∞`, `(1e-7).exp() ≈
  ///    1.0000001`). Either way the bad value reaches a public
  ///    `WordSegment` score, and `align_emissions` returns that
  ///    `WordSegment` straight to the caller — `compose_words` is a
  ///    separate, later, caller-invoked step (`align_emissions` does
  ///    not call it), so whatever it would have sanitised (it maps a
  ///    `NaN` segment score to `0.0` and clamps a positive one back
  ///    into `[0, 1]`, keeping a composed public
  ///    [`Word`](crate::types::Word) NaN-free and in range) is not
  ///    yet applied to a `WordSegment` `align_emissions` has already
  ///    handed back. Both must be rejected here at the seam, not
  ///    left for a downstream step the caller may not even take.
  ///    `−∞` is a plausible caller value for `log(0)` hard-masking,
  ///    but the internal path never produces it and the DP was never
  ///    exercised against it, so it is rejected too.
  ///
  ///    A caller whose own `f32` log-softmax rounds a true `0.0` to a
  ///    tiny positive must clamp (`.min(0.0)`) before constructing.
  ///    `0.0` and `−0.0` are accepted (`log(1) = 0` is a legal
  ///    log-probability); only strictly positive values are not.
  ///
  ///    With `lp ≤ 0` enforced, `exp(lp) ∈ [0, 1]`, and every mean of
  ///    such per-frame values — the per-char mean in `merge_repeats`,
  ///    then the duration-weighted per-word mean in `merge_words` —
  ///    stays in `[0, 1]`. So every `WordSegment` score this pipeline
  ///    produces is finite and in range **by construction**.
  ///
  /// Not `const fn`, unlike most constructors in this crate: the
  /// rejecting paths drop the caller-supplied `data` without moving
  /// it into `Self` (and `Vec<f32>` deallocation isn't
  /// const-evaluable on stable Rust), and the value-domain scan is
  /// not const-evaluable either — an infallible `const fn` could not
  /// reject its input, so matching the crate's usual `pub const fn
  /// new` shape here would mean giving up validation instead.
  ///
  /// # Errors
  ///
  /// Returns [`LogProbsError::Shape`] when `v == 0` or when
  /// `t.checked_mul(v) != Some(data.len())` (product mismatch or
  /// overflow), and [`LogProbsError::Value`] when any element of
  /// `data` is outside the log-probability domain (non-finite, or
  /// finite but `> 0.0`) — reporting the first offending frame,
  /// vocab index, and [`LogProbsValueClass`].
  pub fn new(t: usize, v: usize, data: Vec<f32>) -> Result<Self, LogProbsError> {
    if v == 0 {
      return Err(LogProbsError::Shape(LogProbsShapeError::new(
        t,
        v,
        data.len(),
      )));
    }
    if t.checked_mul(v) != Some(data.len()) {
      return Err(LogProbsError::Shape(LogProbsShapeError::new(
        t,
        v,
        data.len(),
      )));
    }
    // Value-domain scan: every log-probability must be finite AND
    // `≤ 0.0` — exactly the mathematical domain of `log(p)` for
    // `p ∈ (0, 1]`. One short-circuiting O(T×V) pass. It runs only
    // on this external `emissions` seam — the internal ort path
    // builds `LogProbsTV` via a struct literal after
    // `log_softmax_with_finite_guard`, whose output is finite and
    // `≤ 0` by construction (see `new`'s doc), so it never pays this
    // scan. `v > 0` and `data.len() == t * v` hold here, so `idx / v`
    // and `idx % v` recover the (frame, vocab) coordinates. The
    // accept predicate `x.is_finite() && x <= 0.0` rejects `NaN` (all
    // comparisons are false), `±∞` (not finite), and any strictly
    // positive finite; `0.0` and `−0.0` pass (`−0.0 <= 0.0`).
    if let Some(idx) = data.iter().position(|&x| !(x.is_finite() && x <= 0.0)) {
      let bad = data[idx];
      let class = if bad.is_nan() {
        LogProbsValueClass::Nan
      } else if bad.is_infinite() {
        if bad > 0.0 {
          LogProbsValueClass::PosInf
        } else {
          LogProbsValueClass::NegInf
        }
      } else {
        // Finite but failed `x <= 0.0`: strictly positive.
        LogProbsValueClass::Positive
      };
      return Err(LogProbsError::Value(LogProbsValueError::new(
        idx / v,
        idx % v,
        class,
      )));
    }
    Ok(Self { t, v, data })
  }

  /// Build from parts whose contract has ALREADY been established by
  /// the caller — the two internal producers of log-probabilities.
  ///
  /// `pub(crate)`, and it must stay that way. Both callers earn the
  /// skip: `encode_log_softmax` and `Emissions::from_logits` each run
  /// [`log_softmax_with_finite_guard`], whose output is finite and
  /// `<= 0` **by construction** (`lp = (x − max) − ln Σ exp(x − max)`;
  /// the `max` element contributes `exp(0) = 1`, so `ln Σ >= 0` and
  /// every output is `(<= 0) − (>= 0)`). Re-running [`new`](Self::new)'s
  /// `O(T·V)` value-domain scan on that output would be re-deriving a
  /// fact the log-softmax identity already guarantees — and would put a
  /// new per-chunk cost on the ORT hot path, which previously reached
  /// this shape through a bare struct literal.
  ///
  /// This is NOT a public escape hatch, and `Emissions` deliberately has
  /// no equivalent. The seam has exactly one door.
  pub(crate) const fn from_parts_unchecked(t: usize, v: usize, data: Vec<f32>) -> Self {
    Self { t, v, data }
  }

  /// Time dimension (number of wav2vec2 output frames).
  #[must_use]
  pub const fn t(&self) -> usize {
    self.t
  }

  /// Vocab dimension.
  #[must_use]
  pub const fn v(&self) -> usize {
    self.v
  }

  /// Borrow the flat row-major `(T, V)` log-probability buffer.
  #[must_use]
  pub fn data(&self) -> &[f32] {
    &self.data
  }

  /// Read the log-probability of vocab index `v_idx` at frame
  /// `t_idx`, or `None` when `(t_idx, v_idx)` is outside the `(T, V)`
  /// grid.
  ///
  /// The total counterpart of [`at`](Self::at) — the whole `(usize,
  /// usize)` argument space maps to a defined answer, so an
  /// `emissions` caller indexing emissions it did not itself bound can
  /// stay panic-free.
  ///
  /// Bounds-checks `v_idx` against `V` **explicitly**. Checking only
  /// the flat index is not enough: this buffer is row-major, so a
  /// `v_idx >= V` aliases into the *next frame's* row whenever the
  /// flat index still lands inside `data` — `at(0, 3)` on a `(T=2,
  /// V=3)` tensor computes `0 * 3 + 3 = 3`, which is in bounds, and
  /// hands back frame 1's vocab 0.
  #[must_use]
  pub fn get(&self, t_idx: usize, v_idx: usize) -> Option<f32> {
    if v_idx >= self.v {
      return None;
    }
    let idx = t_idx.checked_mul(self.v)?.checked_add(v_idx)?;
    self.data.get(idx).copied()
  }

  /// Read the log-probability of vocab index `v_idx` at frame `t_idx`.
  ///
  /// The pinned DP (`get_trellis` / `backtrack_beam`) indexes through
  /// here on its hot path, always with a `t_idx < T` and a `v_idx < V`
  /// it has already validated (`get_trellis` rejects a `blank_id >= V`
  /// and any token id `>= V` before the first read), so this is
  /// infallible on the `alignment` path and its values are unchanged.
  ///
  /// # Panics
  ///
  /// Panics when `(t_idx, v_idx)` is outside the `(T, V)` grid —
  /// deterministically, with the coordinates in the message, and
  /// **identically in debug and release**. Use [`get`](Self::get) for
  /// the non-panicking form.
  ///
  /// The naked `self.data[t_idx * self.v + v_idx]` this replaces was
  /// neither: on a `(T=2, V=3)` tensor `at(0, 3)` returned frame 1's
  /// vocab 0 — a *silently wrong log-probability*, in both profiles,
  /// because the row-major flat index stayed in bounds — and a large
  /// `t_idx` overflowed the multiply, panicking under debug overflow
  /// checks but wrapping to an in-bounds row in release. Routing
  /// through the checked [`get`](Self::get) closes both: an
  /// out-of-domain coordinate is now always a panic, never a wrong
  /// number. Regressions:
  /// `at_rejects_vocab_index_aliasing_into_the_next_frame`,
  /// `at_rejects_frame_index_that_overflows_the_flat_index`.
  #[must_use]
  pub fn at(&self, t_idx: usize, v_idx: usize) -> f32 {
    match self.get(t_idx, v_idx) {
      Some(lp) => lp,
      None => panic!(
        "LogProbsTV index out of bounds: (t={t_idx}, v={v_idx}) is outside the (T={}, V={}) grid",
        self.t, self.v
      ),
    }
  }
}

/// Run wav2vec2 over `samples_for_aligner` and return per-frame
/// log-probabilities.
///
/// **`samples_for_aligner` must be pre-normalised.** The
/// silence-aware
/// [`crate::runner::aligner::algorithm::normalize::normalize_with_silence_mask`]
/// runs in `Aligner::align` before this function so the silence
/// mask is preserved through preprocessing. Normalisation lives
/// up the call stack so masked regions stay exactly zero in the
/// tensor we feed to ORT.
///
/// The model is expected to take an input named `"input_values"` of
/// shape `(1, T_samples)` and return logits of shape `(1, T_frames,
/// V)`. wav2vec2-base-960h follows this convention; if a different
/// variant uses a different I/O name, parameterise via
/// `Aligner::with_input_name(...)` (not in v1 scope).
///
/// `run_options` carries ONNX Runtime's per-call termination flag;
/// the alignment worker's watchdog calls `RunOptions::terminate()`
/// on timeout, which causes `Session::run_with_options` to surface
/// an error from inside the graph rather than blocking until the
/// model finishes naturally. This is the only way to interrupt a
/// stuck or pathological inference; the `abort_flag` checked at
/// stage boundaries can't help once we are inside `run`.
///
/// Returns `WorkFailure::AlignmentFailed { kind:
/// ModelInferenceFailed, .. }` on any ort error (including a
/// terminate-induced one — the watchdog's
/// `WorkerHangTimeout` is surfaced by the alignment pool wrapper).
///
/// Gated on `feature = "alignment"`: this is the only function in
/// the file that touches `ort`. Everything above and below it
/// (`LogProbsTV`, `log_softmax_with_finite_guard`,
/// `validate_output_dims`, `validate_stride_extent`,
/// `validate_vocab_dim`, `reject_non_finite_input`) is reachable
/// under the ort-free `emissions` feature too.
#[cfg(feature = "alignment")]
pub(crate) fn encode_log_softmax(
  session: &mut Session,
  samples_for_aligner: &[f32],
  run_options: &RunOptions,
  language: &Lang,
) -> Result<LogProbsTV, WorkFailure> {
  let t_samples = samples_for_aligner.len();
  if t_samples == 0 {
    return Err(WorkFailure::Alignment(AlignmentError::ModelInference(
      AlignmentFailure::new(
        SmolStr::from("samples_for_aligner is empty"),
        language.clone(),
      ),
    )));
  }

  // Reject non-finite samples up front with a typed in-band
  // failure. See [`reject_non_finite_input`] for the rationale.
  reject_non_finite_input(samples_for_aligner, language)?;

  // Build a (1, T) f32 input via ort's `(shape, Vec<T>)` tensor
  // constructor — see the module-level NOTE for why we don't go
  // through `ndarray::Array2`. Caller is responsible for the
  // zero-mean / unit-var normalisation (silence-aware variant in
  // `Aligner::align`); the input here goes straight to ORT.
  let input_shape: [i64; 2] = [1, t_samples as i64];
  let input_tensor =
    Tensor::from_array((input_shape, samples_for_aligner.to_vec())).map_err(|e| {
      WorkFailure::Alignment(AlignmentError::ModelInference(AlignmentFailure::new(
        format_smolstr!("Tensor::from_array failed: {e:?}"),
        language.clone(),
      )))
    })?;

  // Most wav2vec2 ONNX exports use the input name "input_values".
  // If the export uses a different name, surface a clear error.
  // `run_with_options` is identical to `run` except it observes
  // the per-call termination flag in `run_options`, so the
  // alignment worker's watchdog can interrupt a stuck graph.
  let outputs = session
    .run_with_options(ort::inputs![input_tensor], run_options)
    .map_err(|e| {
      WorkFailure::Alignment(AlignmentError::ModelInference(AlignmentFailure::new(
        format_smolstr!("Session::run_with_options failed: {e:?}"),
        language.clone(),
      )))
    })?;

  // Take the first (only) output. wav2vec2 has a single logits
  // output; we pull index 0 by name-agnostic iteration.
  let mut iter = outputs.into_iter();
  let (_, output_value) = iter.next().ok_or_else(|| {
    WorkFailure::Alignment(AlignmentError::ModelInference(AlignmentFailure::new(
      SmolStr::from("Session::run returned no outputs"),
      language.clone(),
    )))
  })?;

  let (shape, raw): (&Shape, &[f32]) = output_value.try_extract_tensor::<f32>().map_err(|e| {
    WorkFailure::Alignment(AlignmentError::ModelInference(AlignmentFailure::new(
      format_smolstr!("try_extract_tensor::<f32> failed: {e:?}"),
      language.clone(),
    )))
  })?;

  if shape.len() != 3 || shape[0] != 1 {
    return Err(WorkFailure::Alignment(AlignmentError::ModelInference(
      AlignmentFailure::new(
        format_smolstr!("expected output shape (1, T, V); got {shape:?}"),
        language.clone(),
      ),
    )));
  }

  // Validate the shape integers and their product against the
  // raw buffer length BEFORE we cast to usize / allocate /
  // slice. A model export bug that emits a negative dimension
  // would otherwise wrap to a huge usize and OOM the worker,
  // and a buffer-vs-shape mismatch would panic on the row-slice
  // in `log_softmax_with_finite_guard`.
  let (t, v) = validate_output_dims(shape[1], shape[2], raw.len(), language)?;

  // Validate logits + log-softmax for finiteness. See
  // `log_softmax_with_finite_guard`. The helper returns the neutral
  // `EmissionsError`; the pool path re-maps it to `WorkFailure` here
  // so `encode_log_softmax`'s observable error type is unchanged.
  let data = log_softmax_with_finite_guard(raw, t, v).map_err(|e| e.into_work_failure(language))?;
  // Named constructor for what used to be a bare struct literal. Same
  // cost (no scan): `log_softmax_with_finite_guard`'s output is finite
  // and `<= 0` by construction. See `from_parts_unchecked`.
  Ok(LogProbsTV::from_parts_unchecked(t, v, data))
}

/// Validate the `(T, V)` output dimensions before allocation /
/// slicing.
///
/// Without these checks, a malformed ORT output (negative dim
/// from a buggy export, overflow on `t * v`, or `raw.len()` not
/// matching the declared shape) would either panic the alignment
/// worker on the row-slice, OOM the process on a
/// `Vec::with_capacity` for a wrapped-huge size, or — worst case
/// — silently read into adjacent memory.
///
/// Failure classification:
/// - **Fatal** (`ModelInferenceFailed`) for impossible shapes
/// the operator should hear about: `V <= 0`, negative `T`,
/// `T * V` overflow, or shape-vs-buffer-length mismatch.
/// - **Recoverable** (`NoAlignmentPath`) for `T == 0` with an
/// empty buffer — a chunk shorter than the model's stride
/// produces zero encoder frames; the ASR transcript should
/// surface with `words: []` rather than fail the chunk. A
/// blanket "non-positive dim" rule would turn a data-dependent
/// short-chunk miss into fatal transcript loss.
///
/// Pulled out as a helper so unit tests can drive each branch
/// without an ORT session.
pub(crate) fn validate_output_dims(
  raw_t: i64,
  raw_v: i64,
  raw_len: usize,
  language: &Lang,
) -> Result<(usize, usize), WorkFailure> {
  // V == 0 means the model declared no vocabulary axis — never
  // legitimate. Always fatal.
  if raw_v <= 0 {
    return Err(WorkFailure::Alignment(AlignmentError::ModelInference(
      AlignmentFailure::new(
        format_smolstr!("ORT output has non-positive vocab dim: V={raw_v}"),
        language.clone(),
      ),
    )));
  }
  // Negative T is always a backend bug (truncated /
  // sign-flipped shape descriptor).
  if raw_t < 0 {
    return Err(WorkFailure::Alignment(AlignmentError::ModelInference(
      AlignmentFailure::new(
        format_smolstr!("ORT output has negative time dim: T={raw_t}"),
        language.clone(),
      ),
    )));
  }
  // T == 0 — the model returned a well-formed empty output.
  // With an empty buffer that's a legitimate "chunk too short
  // for any encoder frame" outcome and we surface as
  // recoverable `NoAlignmentPath`. With a non-empty buffer the
  // shape declaration disagrees with the data length — fatal
  // model bug.
  if raw_t == 0 {
    if raw_len != 0 {
      return Err(WorkFailure::Alignment(AlignmentError::ModelInference(
        AlignmentFailure::new(
          format_smolstr!(
            "ORT output declared T=0 but buffer has {raw_len} elements; shape/data mismatch"
          ),
          language.clone(),
        ),
      )));
    }
    return Err(WorkFailure::Alignment(AlignmentError::NoAlignmentPath(
      AlignmentFailure::new(
        SmolStr::from(
          "ORT output has zero encoder frames (chunk too short to align); \
 transcript will surface with words: []",
        ),
        language.clone(),
      ),
    )));
  }
  // i64 → usize is safe after the >0 check on 64-bit; on 32-bit
  // we still want the explicit overflow guard.
  let t = match usize::try_from(raw_t) {
    Ok(v) => v,
    Err(_) => {
      return Err(WorkFailure::Alignment(AlignmentError::ModelInference(
        AlignmentFailure::new(
          format_smolstr!("ORT output T={raw_t} doesn't fit in usize"),
          language.clone(),
        ),
      )));
    }
  };
  let v = match usize::try_from(raw_v) {
    Ok(v) => v,
    Err(_) => {
      return Err(WorkFailure::Alignment(AlignmentError::ModelInference(
        AlignmentFailure::new(
          format_smolstr!("ORT output V={raw_v} doesn't fit in usize"),
          language.clone(),
        ),
      )));
    }
  };
  let total = match t.checked_mul(v) {
    Some(p) => p,
    None => {
      return Err(WorkFailure::Alignment(AlignmentError::ModelInference(
        AlignmentFailure::new(
          format_smolstr!("ORT output dimensions overflow: T={t} * V={v} doesn't fit in usize"),
          language.clone(),
        ),
      )));
    }
  };
  if total != raw_len {
    return Err(WorkFailure::Alignment(AlignmentError::ModelInference(
      AlignmentFailure::new(
        format_smolstr!(
          "ORT output buffer length {raw_len} doesn't match declared T={t} × V={v} = {total}"
        ),
        language.clone(),
      ),
    )));
  }
  Ok((t, v))
}

/// Validate the encoder's frame count against the input audio
/// length. wav2vec2's CNN downsamples by `hop_samples`, so the
/// encoded "time" `T * hop_samples` should lie within
/// `chunk_extent ± 2*hop_samples` (a couple of frames of
/// receptive-field slack on each side).
///
/// Two-sided check — both bounds matter:
///
/// - **Upper bound** (`T * hop > chunk + 2*hop`): the model
/// reports more frames than the input could plausibly support.
/// Either the export uses a smaller stride than `hop_samples`
/// or the configured `hop_samples` is too small. `compose_words`
/// would otherwise emit ranges past the chunk's audio
/// boundary.
/// - **Lower bound** (`T * hop < chunk - 2*hop`): the model
/// reports far fewer frames than the input should produce.
/// Either the export uses a *larger* stride than
/// `hop_samples` or `hop_samples` is too large. `compose_words`
/// would otherwise emit ranges that compress every word into
/// the first portion of the chunk — plausible-looking
/// timestamps that all sit in (e.g.) the first half of the
/// audio.
///
/// `chunk_extent.saturating_sub(slack)` lets very short chunks
/// (where the slack is comparable to `chunk_extent`) pass without
/// false positives — the lower bound clamps to 0. T == 0 cases
/// are already routed to recoverable `NoAlignmentPath` by
/// [`validate_output_dims`].
pub(crate) fn validate_stride_extent(
  t: usize,
  hop_samples: u32,
  chunk_extent: usize,
  language: &Lang,
) -> Result<(), WorkFailure> {
  let frame_extent = (t as u64).saturating_mul(hop_samples as u64);
  let chunk_extent_u64 = chunk_extent as u64;
  let slack = 2u64.saturating_mul(hop_samples as u64);
  let upper_bound = chunk_extent_u64.saturating_add(slack);
  let lower_bound = chunk_extent_u64.saturating_sub(slack);
  if frame_extent > upper_bound {
    return Err(WorkFailure::Alignment(AlignmentError::ModelInference(
      AlignmentFailure::new(
        format_smolstr!(
          "ORT output stride mismatch: T={t} × hop={hop_samples} = {frame_extent} \
 sample-equivalents exceeds chunk ({chunk_extent} samples) + 2-frame slack \
 ({upper_bound}); model export uses a smaller stride than `hop_samples` \
 or `hop_samples` is misconfigured"
        ),
        language.clone(),
      ),
    )));
  }
  if frame_extent < lower_bound {
    return Err(WorkFailure::Alignment(AlignmentError::ModelInference(
      AlignmentFailure::new(
        format_smolstr!(
          "ORT output stride mismatch: T={t} × hop={hop_samples} = {frame_extent} \
 sample-equivalents below chunk ({chunk_extent} samples) − 2-frame slack \
 ({lower_bound}); model export uses a larger stride than `hop_samples` \
 or `hop_samples` is misconfigured"
        ),
        language.clone(),
      ),
    )));
  }
  Ok(())
}

/// Validate the model output's vocab dimension against the
/// tokenizer's vocab size. A wrong ONNX export (e.g. a hidden-
/// states tensor with a much larger trailing dim, or a CTC head
/// trained on a different alphabet) would otherwise pass the
/// per-token id check in `ctc_viterbi` whenever the chunk's
/// in-vocab token ids happen to fit, then read posteriors from
/// columns the tokenizer thinks correspond to the wrong tokens
/// — emitting believable but corrupt timings.
///
/// Strict equality matches the wav2vec2 ASR family (model output
/// dim == tokenizer vocab size, including special tokens like
/// `<pad>` / `<s>` / `</s>` / `<unk>` / `|`). If a future
/// downstream model legitimately has a different output width,
/// this helper would need a configured override; for the
/// supported family, mismatch is always a model/tokenizer
/// pairing bug.
pub(crate) fn validate_vocab_dim(
  v: usize,
  expected_v: usize,
  language: &Lang,
) -> Result<(), WorkFailure> {
  if v != expected_v {
    return Err(WorkFailure::Alignment(AlignmentError::ModelInference(
      AlignmentFailure::new(
        format_smolstr!(
          "ORT output vocab dim V={v} doesn't match tokenizer vocab size {expected_v}; \
 model and tokenizer are paired incorrectly — Viterbi would otherwise read \
 posteriors from columns that don't correspond to the tokenizer's tokens"
        ),
        language.clone(),
      ),
    )));
  }
  Ok(())
}

/// Compute row-major `(T, V)` log-softmax of `raw` with a fatal
/// finiteness guard.
///
/// Without this guard a NaN / ±inf in any logit produced a NaN
/// row in the output; Viterbi then computed a non-finite final
/// `dp_prev` and surfaced `NoAlignmentPath`, which — on the
/// `alignment`-feature pool path — the alignment pool classifies as
/// recoverable, silently swallowing a backend numeric failure
/// (model export bug, GPU / ORT regression, NaN propagation from
/// upstream) as "no words". This helper checks each row's input and
/// the resulting `log_z`; either non-finite returns
/// [`EmissionsError::Numeric`] instead. On the `alignment` pool path
/// the orchestrator re-maps that into `WorkFailure` and the runner
/// emits `Event::Error` so the operator learns about the broken
/// backend; an `emissions`-only caller (no pool, no runner) gets the
/// neutral typed error back directly as this function's `Result::Err`.
///
/// Pulled out of `encode_log_softmax`'s body so unit tests can
/// exercise the rejection paths without a `Session`, and so
/// callers under the ort-free `emissions` feature (their own
/// encoder's raw logits, no ort session at all) can run the same
/// finiteness-guarded log-softmax asry's own `alignment` path
/// uses. `pub` for exactly that reason — reachable at
/// `asry::emissions::log_softmax_with_finite_guard`.
///
/// Also validates `v != 0` and `t.checked_mul(v) ==
/// Some(raw.len())` up front, before any indexing into `raw` —
/// overflow-safe, the same discipline `validate_output_dims`
/// applies to the ORT-sourced shape. Internal callers already get
/// `t`/`v` pre-validated against the ORT output there (redundant,
/// O(1), here); external `emissions`-feature callers construct
/// `t`/`v`/`raw` themselves and are not.
///
/// # Errors
///
/// Returns [`EmissionsError::Shape`] when `v == 0` (checked first,
/// for every `t` including `t == 0`: a CTC vocabulary must contain at
/// least the blank token, so a zero-length vocab axis is never valid)
/// or when `t.checked_mul(v) != Some(raw.len())` (product mismatch or
/// `usize` overflow); returns [`EmissionsError::Numeric`] when any
/// caller logit or resulting log-probability is non-finite.
pub fn log_softmax_with_finite_guard(
  raw: &[f32],
  t: usize,
  v: usize,
) -> Result<Vec<f32>, EmissionsError> {
  // V == 0 means the caller declared no vocabulary axis — never
  // legitimate, for any T (including T == 0, which would
  // otherwise skip the frame loop below entirely and return
  // `Ok(vec![])` without ever looking at V). Checked first, the
  // same discipline `validate_output_dims` applies to the
  // ORT-sourced shape. Reuses the `LogProbsTV::new` shape leaf so the
  // seam reports one dimension/product/overflow vocabulary.
  if v == 0 {
    return Err(EmissionsError::Shape(LogProbsShapeError::new(
      t,
      v,
      raw.len(),
    )));
  }
  let Some(total) = t.checked_mul(v) else {
    return Err(EmissionsError::Shape(LogProbsShapeError::new(
      t,
      v,
      raw.len(),
    )));
  };
  if total != raw.len() {
    return Err(EmissionsError::Shape(LogProbsShapeError::new(
      t,
      v,
      raw.len(),
    )));
  }
  let mut data = Vec::with_capacity(total);
  for t_idx in 0..t {
    let row = &raw[t_idx * v..(t_idx + 1) * v];
    if let Some(bad_v) = row.iter().position(|x| !x.is_finite()) {
      return Err(EmissionsError::Numeric(EmissionsFailure::new(
        format_smolstr!(
          "encoder supplied non-finite logit at frame {t_idx}, vocab {bad_v}: {}",
          row[bad_v]
        ),
      )));
    }
    // Shifted log-sum-exp computed in f64. We do NOT add `max`
    // back into f32 to form `log_z` and then subtract it again,
    // because for a row with a large common offset (e.g.
    // `[1e20, 1e20]`) `sum.ln()` rounds away when added to `max`
    // in f32 — `max + sum.ln() as f32 = 1e20 + 0.69 ≈ 1e20` in
    // f32 — and every `lp = x - log_z` then collapses to `0.0`
    // instead of the correct `-ln(2)`. The output passes the
    // finiteness checks but is no longer a log-probability,
    // hiding the backend numeric skew as plausible-looking
    // alignment input. Flagged this; the fix is
    // to keep the subtraction of `max` in shifted f64 space and
    // only cast the final `lp` to f32 (where `lp = (x - max) -
    // sum.ln()` is bounded between `-inf..=0` and never needs
    // `max` to fold in).
    let max = row.iter().copied().fold(f32::NEG_INFINITY, f32::max);
    let max_f64 = max as f64;
    let mut sum = 0.0_f64;
    for &x in row {
      sum += ((x as f64) - max_f64).exp();
    }
    let log_z_shifted = sum.ln();
    if !log_z_shifted.is_finite() {
      // `sum == 0.0` (whole row was -inf, or every shifted exp
      // underflowed) → ln(0) = -inf. `sum < 0` is impossible
      // here. `sum.ln() == NaN` is impossible from non-negative
      // f64 input. So this branch fires only on the all-(-inf)
      // case, which the existing all-`NEG_INFINITY` regression
      // also covers.
      return Err(EmissionsError::Numeric(EmissionsFailure::new(
        format_smolstr!(
          "log-softmax shifted normaliser non-finite at frame {t_idx}: \
 sum.ln()={log_z_shifted}, max={max}"
        ),
      )));
    }
    // Per-output log-probability. `lp_f64 = (x - max) - sum.ln()`
    // is bounded in `(-∞, 0]` for any finite input row (since
    // `(x - max) <= 0` and `sum.ln() >= 0` whenever any row
    // element equals `max`). We still keep the per-element
    // finiteness check because pathological inputs like
    // `[f32::MAX, -f32::MAX]` can underflow `lp_f64 as f32` to
    // `-inf`; surfacing that as `ModelInferenceFailed` keeps the
    // backend-numeric-failure path typed (a `-inf` slipping into
    // `data` would be visible to Viterbi as a valid-but-very-low
    // log-prob, masking the bug as `NoAlignmentPath` /
    // `words: []`).
    for &x in row {
      let lp_f64 = ((x as f64) - max_f64) - log_z_shifted;
      let lp = lp_f64 as f32;
      if !lp.is_finite() {
        return Err(EmissionsError::Numeric(EmissionsFailure::new(
          format_smolstr!(
            "log-softmax output non-finite at frame {t_idx}: \
 x={x}, max={max}, sum_ln={log_z_shifted}, lp={lp}"
          ),
        )));
      }
      data.push(lp);
    }
  }
  Ok(data)
}

/// Reject non-finite (NaN / ±inf) samples before any audio
/// processing runs.
///
/// Without this guard, a single bad sample propagates through
/// `zero_mean_unit_var_normalize`'s mean/variance reductions
/// (NaN poisons every downstream f64 op) and ends up in the
/// tensor we hand to ORT. The model then returns either NaN
/// logits (every word gets a NaN score) or the chunk fails
/// downstream as `NoAlignmentPath` with no clue why.
///
/// Pulled out as a helper so the unit tests can exercise the
/// rejection path without spinning up a `Session` (the public
/// `encode_log_softmax` consumes one). The `Aligner::align`
/// integration tests cover the full encode path against the
/// real ORT fixture.
pub(crate) fn reject_non_finite_input(samples: &[f32], language: &Lang) -> Result<(), WorkFailure> {
  if let Some(bad_idx) = samples.iter().position(|s| !s.is_finite()) {
    return Err(WorkFailure::Alignment(AlignmentError::ModelInference(
      AlignmentFailure::new(
        format_smolstr!(
          "samples_for_aligner contains non-finite value at index {bad_idx}: {}",
          samples[bad_idx]
        ),
        language.clone(),
      ),
    )));
  }
  Ok(())
}

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

  /// Pure log-softmax math sanity check. Doesn't touch ort.
  #[test]
  fn log_softmax_sums_to_zero_in_log_space() {
    let row = [1.0f32, 2.0, 3.0];
    let max = row.iter().copied().fold(f32::NEG_INFINITY, f32::max);
    let mut sum = 0.0_f64;
    for &x in &row {
      sum += ((x - max) as f64).exp();
    }
    let log_z = max + (sum.ln() as f32);
    let lp: Vec<f32> = row.iter().map(|x| x - log_z).collect();
    let exp_sum: f32 = lp.iter().map(|x| x.exp()).sum();
    assert!((exp_sum - 1.0).abs() < 1e-5, "softmax must sum to 1");
    for &v in &lp {
      assert!(v <= 0.0, "log-prob must be <= 0");
    }
  }

  /// Regression: NaN / ±inf input must fail in-band with
  /// `ModelInferenceFailed` before the scalar normaliser runs.
  /// The error message names the offending index so a
  /// downstream operator has a hook for debugging upstream audio
  /// pipelines.
  #[test]
  fn reject_non_finite_input_flags_nan() {
    use crate::types::Lang;
    let samples = vec![0.1_f32, 0.2, f32::NAN, 0.4];
    let err = reject_non_finite_input(&samples, &Lang::En).unwrap_err();
    match err {
      WorkFailure::Alignment(AlignmentError::ModelInference(payload)) => {
        assert!(
          payload.message().contains("index 2"),
          "message must name index; got {message}",
          message = payload.message()
        );
      }
      other => panic!("expected AlignmentFailed; got {other:?}"),
    }
  }

  #[test]
  fn reject_non_finite_input_flags_positive_infinity() {
    use crate::types::Lang;
    let samples = vec![0.0_f32, f32::INFINITY];
    assert!(reject_non_finite_input(&samples, &Lang::En).is_err());
  }

  #[test]
  fn reject_non_finite_input_flags_negative_infinity() {
    use crate::types::Lang;
    let samples = vec![f32::NEG_INFINITY, 0.0_f32];
    assert!(reject_non_finite_input(&samples, &Lang::En).is_err());
  }

  #[test]
  fn reject_non_finite_input_passes_finite_audio() {
    use crate::types::Lang;
    // Both ordinary [-1, 1] audio and high-magnitude finite
    // inputs are accepted at this layer — magnitude precision
    // is the SIMD-precision-guard's job, not this guard's.
    let samples = vec![-1.0_f32, 0.0, 1.0, 1e10, -1e10];
    assert!(reject_non_finite_input(&samples, &Lang::En).is_ok());
  }

  // -------- LogProbsTV::new validation (the `emissions` seam) --------

  /// `LogProbsTV::new` is the public, validating constructor a
  /// caller with its own encoder (no ort session, `emissions`
  /// feature) uses to hand its output to the alignment pipeline.
  /// A correctly-shaped buffer must construct successfully and
  /// preserve the dimensions + data verbatim.
  #[test]
  fn new_accepts_matching_shape() {
    let lp = LogProbsTV::new(2, 3, vec![-1.0, -2.0, -3.0, -4.0, -5.0, -6.0]).expect("2 * 3 == 6");
    assert_eq!(lp.t(), 2);
    assert_eq!(lp.v(), 3);
    assert_eq!(lp.data(), &[-1.0, -2.0, -3.0, -4.0, -5.0, -6.0]);
  }

  /// A buffer shorter than `t * v` must be rejected rather than
  /// silently accepted (which would let `at()` panic on an
  /// out-of-bounds index later, deep inside the trellis loop,
  /// instead of at the construction boundary).
  ///
  /// `let Err(err) = ... else { panic!() }` rather than
  /// `.unwrap_err()`: `LogProbsTV` (the `Ok` side) intentionally
  /// carries no `Debug` impl — its `data: Vec<f32>` can be a
  /// wav2vec2-scale emission matrix, not something worth printing
  /// wholesale on a mismatched test expectation — and
  /// `Result::unwrap_err` requires `T: Debug`.
  #[test]
  fn new_rejects_undersized_buffer() {
    let Err(LogProbsError::Shape(err)) = LogProbsTV::new(2, 3, vec![0.0_f32; 5]) else {
      panic!("t=2, v=3, data.len()=5 must be rejected as a shape mismatch");
    };
    assert_eq!(err.t(), 2);
    assert_eq!(err.v(), 3);
    assert_eq!(err.data_len(), 5);
    assert!(err.to_string().contains("t=2"));
    assert!(err.to_string().contains("v=3"));
    assert!(err.to_string().contains("data.len()=5"));
  }

  /// A buffer longer than `t * v` is equally a shape mismatch —
  /// `data.len()` must equal the product exactly, not merely be
  /// large enough.
  #[test]
  fn new_rejects_oversized_buffer() {
    let Err(LogProbsError::Shape(err)) = LogProbsTV::new(2, 3, vec![0.0_f32; 7]) else {
      panic!("t=2, v=3, data.len()=7 must be rejected as a shape mismatch");
    };
    assert_eq!(err.data_len(), 7);
  }

  /// `t * v` overflowing `usize` must reject via `checked_mul`
  /// rather than silently wrap to a small product that happens to
  /// equal a small `data.len()` — the same overflow discipline
  /// `validate_output_dims` already applies to the ORT-sourced
  /// shape.
  #[test]
  fn new_rejects_t_v_overflow() {
    let big = usize::MAX / 2 + 1;
    let Err(LogProbsError::Shape(err)) = LogProbsTV::new(big, big, Vec::new()) else {
      panic!("t * v overflowing usize must be rejected, not silently wrapped");
    };
    assert_eq!(err.t(), big);
    assert_eq!(err.v(), big);
    assert_eq!(err.data_len(), 0);
  }

  /// `t.checked_mul(v) == Some(data.len())` holds trivially for
  /// `t=0, v=0, data=[]` (`0 * 0 == 0`), so the shape check alone
  /// used to let a zero-length vocab axis through — a `LogProbsTV`
  /// with no vocabulary at all, not even a blank token, is
  /// meaningless for CTC. Must reject regardless of `t`; `t=0` is
  /// the degenerate case that used to slip past the shape check
  /// entirely (see `validate_output_dims`'s identical `V == 0`
  /// rule for the ORT-sourced counterpart of this seam).
  #[test]
  fn new_rejects_zero_vocab_with_zero_t() {
    let Err(LogProbsError::Shape(err)) = LogProbsTV::new(0, 0, Vec::new()) else {
      panic!("t=0, v=0 must be rejected: a CTC vocabulary needs at least the blank token");
    };
    assert_eq!(err.t(), 0);
    assert_eq!(err.v(), 0);
    assert_eq!(err.data_len(), 0);
    let message = err.to_string();
    assert!(
      message.contains("zero-length vocab"),
      "message must call out the zero-length vocab dim, not a shape \
 mismatch (t * v == data.len() actually holds here); got {message}"
    );
    assert!(
      !message.contains("expected data.len()"),
      "must not reuse the shape-mismatch wording, which would falsely \
 claim t * v != data.len(); got {message}"
    );
  }

  /// Same defect, `t=1`: `1 * 0 == 0 == data.len()`, so the shape
  /// check alone accepts it too. Distinct from `t=0` because `at()`
  /// would be reachable (frames exist) with no vocab columns to
  /// index — confirms the rejection isn't a `T == 0` special case.
  #[test]
  fn new_rejects_zero_vocab_with_positive_t() {
    let Err(LogProbsError::Shape(err)) = LogProbsTV::new(1, 0, Vec::new()) else {
      panic!("t=1, v=0 must be rejected: a CTC vocabulary needs at least the blank token");
    };
    assert_eq!(err.t(), 1);
    assert_eq!(err.v(), 0);
    assert_eq!(err.data_len(), 0);
    assert!(err.to_string().contains("zero-length vocab"));
  }

  /// Regression (codex round 4): the exact failing history —
  /// `LogProbsTV::new(1, 2, vec![f32::NAN, 0.0])`, `blank_id = 0` —
  /// reached the DP with a `NaN` emission in the blank column. The
  /// one-cell trellis end-cell (a *different* cell) stayed finite
  /// and passed the backtrack guard, so the seed confidence
  /// `at(final_t, blank).exp()` = `NaN.exp()` = `NaN` propagated
  /// into a public `Word`, violating its `[0, 1]` NaN-free score
  /// contract. `new` now rejects the non-finite emission up front:
  /// a typed `Err`, never `Ok`.
  #[test]
  fn new_rejects_nan_from_codex_failing_history() {
    let Err(LogProbsError::Value(err)) = LogProbsTV::new(1, 2, vec![f32::NAN, 0.0]) else {
      panic!("a NaN emission must be rejected as a value-domain error, never accepted");
    };
    assert_eq!(err.frame(), 0);
    assert_eq!(err.vocab_index(), 0);
    assert_eq!(err.class(), LogProbsValueClass::Nan);
  }

  /// `+∞` swept into the token column (vocab 1) of a multi-frame
  /// lattice is rejected as a value-domain error locating the exact
  /// cell and class.
  #[test]
  fn new_rejects_positive_infinity_in_token_column() {
    // 2 frames × 2 vocab; +inf at frame 1, vocab 1.
    let mut data = vec![-1.0_f32, -2.0, -3.0, -4.0];
    data[3] = f32::INFINITY;
    let Err(LogProbsError::Value(err)) = LogProbsTV::new(2, 2, data) else {
      panic!("a +inf emission must be rejected");
    };
    assert_eq!(err.frame(), 1);
    assert_eq!(err.vocab_index(), 1);
    assert_eq!(err.class(), LogProbsValueClass::PosInf);
  }

  /// `−∞` — a plausible `log(0)` hard-mask value a caller might
  /// pass — is rejected too: the internal path never produces it
  /// and the DP was never exercised against it, so the seam keeps
  /// its domain equal to the tested (all-finite) one. Placed in the
  /// blank column of the final frame.
  #[test]
  fn new_rejects_negative_infinity_hard_mask_value() {
    // 2 frames × 2 vocab; -inf at frame 1, vocab 0 (blank column).
    let mut data = vec![-1.0_f32, -2.0, -3.0, -4.0];
    data[2] = f32::NEG_INFINITY;
    let Err(LogProbsError::Value(err)) = LogProbsTV::new(2, 2, data) else {
      panic!("a -inf emission must be rejected");
    };
    assert_eq!(err.frame(), 1);
    assert_eq!(err.vocab_index(), 0);
    assert_eq!(err.class(), LogProbsValueClass::NegInf);
  }

  /// A `NaN` specifically in the FINAL frame's blank column — the
  /// cell the backtrack seed reads (`at(final_t, blank).exp()`), the
  /// exact bypass the shape and end-cell guards missed — is
  /// rejected with precise coordinates.
  #[test]
  fn new_rejects_nan_in_final_frame_blank_column() {
    // 3 frames × 2 vocab; NaN at final frame (2), blank (0).
    let mut data = vec![-1.0_f32; 6];
    data[4] = f32::NAN;
    let Err(LogProbsError::Value(err)) = LogProbsTV::new(3, 2, data) else {
      panic!("a NaN in the final-frame blank column must be rejected");
    };
    assert_eq!(err.frame(), 2);
    assert_eq!(err.vocab_index(), 0);
    assert_eq!(err.class(), LogProbsValueClass::Nan);
  }

  /// Round-5 domain tightening (was round 4's "positivity
  /// unvalidated" accept-test, now inverted — that carve-out was
  /// wrong). A finite positive value is not a log-probability
  /// (`log(p) ≤ 0` for `p ∈ (0, 1]`) and is rejected. A tiny
  /// positive `1e-7` — the exact shape a caller's `f32` log-softmax
  /// rounding a true `0.0` upward produces — exponentiates to
  /// `≈ 1.0000001`, out of `[0, 1]`, and reaches a public
  /// `WordSegment` score before any clamp can defend it; the caller
  /// must clamp with `.min(0.0)` before constructing.
  #[test]
  fn new_rejects_tiny_positive_value() {
    let Err(LogProbsError::Value(err)) = LogProbsTV::new(1, 2, vec![1.0e-7_f32, -0.5]) else {
      panic!("a finite positive value must be rejected as out-of-domain");
    };
    assert_eq!(err.frame(), 0);
    assert_eq!(err.vocab_index(), 0);
    assert_eq!(err.class(), LogProbsValueClass::Positive);
  }

  /// Regression (codex round 5): the exact failing history —
  /// `LogProbsTV::new(1, 2, vec![f32::MAX, -1.0])`, `blank_id = 0` —
  /// passed round 4's finite-only scan (`f32::MAX` is finite) and
  /// reached the DP, where the final-frame blank seed
  /// `at(final_t, blank).exp() = f32::MAX.exp() = +∞` propagated into
  /// a public `WordSegment` score (empirically verified `+∞` before
  /// this fix). `f32::MAX` is finite but `> 0`, so `new` now rejects
  /// it as an out-of-domain value at the first offending cell.
  #[test]
  fn new_rejects_f32_max_from_codex_failing_history() {
    let Err(LogProbsError::Value(err)) = LogProbsTV::new(1, 2, vec![f32::MAX, -1.0]) else {
      panic!("f32::MAX (finite but > 0) must be rejected as out-of-domain");
    };
    assert_eq!(err.frame(), 0);
    assert_eq!(err.vocab_index(), 0);
    assert_eq!(err.class(), LogProbsValueClass::Positive);
  }

  /// The `≤ 0` bound is inclusive: `0.0` (`log(1)`, the maximal
  /// log-probability) and `−0.0` are legal and must be accepted
  /// (`−0.0 <= 0.0` holds). Only strictly positive finite values are
  /// rejected.
  #[test]
  fn new_accepts_zero_and_negative_zero() {
    let lp =
      LogProbsTV::new(1, 3, vec![0.0_f32, -0.0, -1.0]).expect("0.0 and -0.0 are ≤ 0, so accepted");
    assert_eq!(lp.t(), 1);
    assert_eq!(lp.v(), 3);
    assert_eq!(lp.at(0, 0), 0.0);
    assert_eq!(lp.at(0, 1), -0.0);
  }

  /// `LogProbsError`'s `Display` is transparent — each arm forwards
  /// verbatim to the wrapped shape / value error
  /// (`#[error(transparent)]`).
  #[test]
  fn log_probs_error_display_is_transparent() {
    let Err(shape) = LogProbsTV::new(2, 3, vec![0.0_f32; 5]) else {
      panic!("shape mismatch expected");
    };
    let LogProbsError::Shape(inner) = shape else {
      panic!("expected the Shape arm");
    };
    assert_eq!(LogProbsError::Shape(inner).to_string(), inner.to_string());

    let Err(value) = LogProbsTV::new(1, 1, vec![f32::NAN]) else {
      panic!("value error expected");
    };
    let LogProbsError::Value(inner) = value else {
      panic!("expected the Value arm");
    };
    assert_eq!(LogProbsError::Value(inner).to_string(), inner.to_string());
    assert!(inner.to_string().contains("out of domain"));
  }

  /// `t=0` with a *positive* vocab dim is the legitimate degenerate
  /// case (a chunk too short to produce any encoder frame) and must
  /// stay `Ok` — the zero-vocab rejection above must not overreach
  /// into rejecting `T == 0` itself.
  #[test]
  fn new_accepts_zero_t_with_positive_vocab() {
    let lp = LogProbsTV::new(0, 5, Vec::new()).expect("t=0 with v=5 and an empty buffer is valid");
    assert_eq!(lp.t(), 0);
    assert_eq!(lp.v(), 5);
    assert!(lp.data().is_empty());
  }

  #[test]
  fn at_indexes_correctly() {
    let lp = LogProbsTV {
      t: 2,
      v: 3,
      data: vec![-1.0, -2.0, -3.0, -4.0, -5.0, -6.0],
    };
    assert_eq!(lp.at(0, 0), -1.0);
    assert_eq!(lp.at(0, 2), -3.0);
    assert_eq!(lp.at(1, 0), -4.0);
    assert_eq!(lp.at(1, 2), -6.0);
  }

  /// Helper for the `at` / `get` domain tests: the same well-formed
  /// `(T=2, V=3)` tensor `at_indexes_correctly` pins, built through the
  /// validating public constructor so the tests speak to the seam an
  /// `emissions` caller actually reaches.
  fn two_by_three() -> LogProbsTV {
    LogProbsTV::new(2, 3, vec![-1.0, -2.0, -3.0, -4.0, -5.0, -6.0])
      .expect("2 * 3 == 6 and every value is a valid log-probability")
  }

  /// `get` is the total form of `at`: every `(usize, usize)` maps to a
  /// defined answer, `Some` inside the grid and `None` outside it —
  /// including the two coordinates that used to alias or overflow.
  #[test]
  fn get_is_total_over_its_argument_space() {
    let lp = two_by_three();
    assert_eq!(lp.get(0, 0), Some(-1.0));
    assert_eq!(lp.get(1, 2), Some(-6.0));
    // v_idx == V: out of the grid even though the flat index is in
    // bounds (this is the aliasing case).
    assert_eq!(lp.get(0, 3), None);
    // t_idx == T: past the last row.
    assert_eq!(lp.get(2, 0), None);
    // Flat index overflows `usize`.
    assert_eq!(lp.get(usize::MAX, 0), None);
    assert_eq!(lp.get(usize::MAX, usize::MAX), None);
  }

  ///  class regression (the `build_speech_frames` sibling).
  /// `LogProbsTV::at` is `pub` on the `emissions` surface, so a caller
  /// can hand it any `(usize, usize)`. The naked
  /// `self.data[t_idx * self.v + v_idx]` did not close that domain.
  ///
  /// The buffer is ROW-MAJOR, so a `v_idx >= V` walks into the *next
  /// frame's* row whenever the flat index still lands inside `data`.
  /// On this `(T=2, V=3)` tensor, `at(0, 3)` computed `0 * 3 + 3 = 3`
  /// — in bounds — and returned `-4.0`, which is frame 1's vocab 0.
  /// Not a panic: a **silently wrong log-probability, in debug and
  /// release alike**, fed straight into a CTC emission.
  ///
  /// It must now be a panic, not a wrong number.
  #[test]
  #[should_panic(expected = "(t=0, v=3) is outside the (T=2, V=3) grid")]
  fn at_rejects_vocab_index_aliasing_into_the_next_frame() {
    let lp = two_by_three();
    // Pre-fix: returns -4.0 (frame 1, vocab 0). Post-fix: panics.
    let _ = lp.at(0, 3);
  }

  /// The overflow half of the same hole: a large `t_idx` wraps
  /// `t_idx * self.v`. Under debug overflow checks that panicked with
  /// the opaque "attempt to multiply with overflow"; in release it
  /// wrapped to a small in-bounds flat index and returned an unrelated
  /// element. `checked_mul` in `get` makes both profiles agree on a
  /// deterministic, coordinate-naming panic.
  #[test]
  #[should_panic(expected = "is outside the (T=2, V=3) grid")]
  fn at_rejects_frame_index_that_overflows_the_flat_index() {
    let lp = two_by_three();
    // (usize::MAX / 3) + 1 multiplied by V=3 wraps to 1 mod 2^64, so
    // pre-fix release handed back `data[1]`.
    let _ = lp.at((usize::MAX / 3) + 1, 0);
  }

  /// `at` stays infallible on every coordinate the pinned DP produces:
  /// `get_trellis` validates `blank_id < V` and every token id `< V`
  /// before its first read, and only ever reads rows `< T`. Sweep the
  /// whole grid to pin that the hardened `at` agrees with `get` and
  /// with the flat row-major layout — i.e. the `alignment` path's
  /// values are unchanged.
  #[test]
  fn at_agrees_with_get_across_the_whole_grid() {
    let lp = two_by_three();
    for t_idx in 0..lp.t() {
      for v_idx in 0..lp.v() {
        let expected = lp.data()[t_idx * lp.v() + v_idx];
        assert_eq!(lp.at(t_idx, v_idx), expected);
        assert_eq!(lp.get(t_idx, v_idx), Some(expected));
      }
    }
  }

  /// NaN logits from a broken backend must surface as a fatal
  /// numeric failure, not get swallowed into NaN log-probs that
  /// Viterbi later classifies as `NoAlignmentPath` (the recoverable
  /// bucket). Under `default-features = false, features =
  /// ["emissions"]` `ort` isn't even linked and the input came from
  /// the caller's own encoder, so the message must not attribute the
  /// failure to ORT specifically — it names the generic `encoder`.
  #[test]
  fn log_softmax_rejects_nan_logits_with_numeric_failure() {
    let raw = vec![0.0_f32, f32::NAN, 0.0]; // 1×3
    let err = log_softmax_with_finite_guard(&raw, 1, 3).unwrap_err();
    match err {
      EmissionsError::Numeric(payload) => {
        let message = payload.message();
        assert!(
          message.contains("non-finite logit"),
          "message must call out the non-finite logit; got {message}"
        );
        assert!(message.contains("frame 0"));
        assert!(message.contains("vocab 1"));
        assert!(
          !message.to_ascii_lowercase().contains("ort"),
          "message must not attribute the failure to ORT specifically — \
`log_softmax_with_finite_guard` is also reachable ort-free under the \
`emissions` feature, where the non-finite value came from the caller's \
own encoder; got {message:?}"
        );
      }
      other => panic!("expected EmissionsError::Numeric; got {other:?}"),
    }
  }

  #[test]
  fn log_softmax_rejects_positive_infinity_logits() {
    let raw = vec![0.0_f32, f32::INFINITY, 0.0];
    assert!(log_softmax_with_finite_guard(&raw, 1, 3).is_err());
  }

  #[test]
  fn log_softmax_rejects_negative_infinity_logits() {
    let raw = vec![f32::NEG_INFINITY, 0.0, 0.0];
    assert!(log_softmax_with_finite_guard(&raw, 1, 3).is_err());
  }

  /// `f32::NEG_INFINITY.is_finite()` returns false, so an all-`-inf`
  /// row fails at the per-element finiteness check and surfaces as a
  /// numeric failure rather than a NaN row Viterbi would misread.
  #[test]
  fn log_softmax_rejects_all_neg_infinity_row() {
    let raw = vec![f32::NEG_INFINITY; 3];
    assert!(log_softmax_with_finite_guard(&raw, 1, 3).is_err());
  }

  /// Finite extreme logits can produce a non-finite output
  /// log-prob: `[f32::MAX, -f32::MAX]` has finite input, finite max,
  /// finite log_z (= f32::MAX), but the second element's `x - log_z
  /// = -f32::MAX - f32::MAX = -inf`. The `-inf` would be stored in
  /// `data`; Viterbi would later return `NoAlignmentPath`
  /// (recoverable), hiding a real numeric failure as `words: []`.
  /// The per-element finite check surfaces it as a numeric failure.
  #[test]
  fn log_softmax_rejects_finite_extremes_that_overflow_lp() {
    let raw = vec![f32::MAX, -f32::MAX];
    let err = log_softmax_with_finite_guard(&raw, 1, 2).unwrap_err();
    let EmissionsError::Numeric(payload) = err else {
      panic!("expected EmissionsError::Numeric; got {err:?}");
    };
    let message = payload.message();
    assert!(
      message.contains("log-softmax output non-finite"),
      "diagnostic must call out the per-element finite check; got {message}",
    );
  }

  /// Sanity: a finite, well-behaved row produces a finite
  /// log-softmax row that sums to 1 in linear space.
  #[test]
  fn log_softmax_finite_input_roundtrips() {
    let raw = vec![1.0_f32, 2.0, 3.0];
    let out = log_softmax_with_finite_guard(&raw, 1, 3).expect("ok");
    assert_eq!(out.len(), 3);
    assert!(out.iter().all(|x| x.is_finite()));
    let sum: f32 = out.iter().map(|x| x.exp()).sum();
    assert!((sum - 1.0).abs() < 1e-5);
  }

  #[test]
  fn log_softmax_large_common_offset_normalises_to_unit_exp_sum() {
    // regression: the previous implementation
    // computed `log_z = max + sum.ln() as f32`. For a row with
    // a large common offset like `[1e20, 1e20]`, `sum.ln() =
    // ln(2) ≈ 0.69` rounds away when added to `max = 1e20` in
    // f32 — `1e20 + 0.69 ≈ 1e20` — so each `lp = x - log_z`
    // collapsed to `0.0` instead of the correct `-ln(2)`. The
    // outputs passed the finiteness check but were no longer
    // log-probabilities, hiding a backend numeric failure as
    // plausible alignment output. The fix keeps the
    // subtraction of `max` in f64 (`lp_f64 = (x as f64 -
    // max_f64) - sum.ln()`) so the shifted log-prob is correct
    // regardless of `max`'s magnitude.
    let raw = vec![1.0e20_f32, 1.0e20_f32];
    let out = log_softmax_with_finite_guard(&raw, 1, 2).expect("ok");
    assert_eq!(out.len(), 2);
    assert!(out.iter().all(|x| x.is_finite()));
    // Exp-sum should be 1.0 (i.e. softmax probabilities sum to 1).
    let exp_sum: f32 = out.iter().map(|x| x.exp()).sum();
    assert!(
      (exp_sum - 1.0).abs() < 1e-5,
      "exp(lp) sum must equal 1, got {exp_sum}; lps = {out:?}"
    );
    // Each lp should be log(0.5) = -ln(2) ≈ -0.6931.
    for lp in &out {
      assert!(
        (lp - (-(2.0_f32).ln())).abs() < 1e-4,
        "expected ~{}, got {lp}",
        -(2.0_f32).ln()
      );
    }
  }

  // --- log_softmax_with_finite_guard dimension validation ---
  //
  // `encode_log_softmax`'s internal call is already preceded by
  // `validate_output_dims`, so these paths are unreachable from
  // that caller. They matter for the `emissions`-feature entry
  // point, which lets an external caller supply `t`/`v`/`raw`
  // directly with no ORT-side pre-validation at all.

  /// `raw` shorter than the declared `T × V` used to index straight
  /// past the end of the last row's slice — an unchecked-slicing
  /// panic instead of a typed error. Now a typed shape rejection.
  #[test]
  fn log_softmax_rejects_undersized_buffer() {
    let err = log_softmax_with_finite_guard(&[0.0f32, 0.0], 2, 2).unwrap_err();
    let EmissionsError::Shape(payload) = err else {
      panic!("expected EmissionsError::Shape; got {err:?}");
    };
    assert!(
      payload.to_string().contains("shape mismatch"),
      "diagnostic must call out the shape mismatch; got {payload}",
    );
  }

  /// An empty buffer against non-zero declared dims is the same
  /// unchecked-slicing panic as the general undersized case, called
  /// out on its own since an empty `raw` is the most likely
  /// caller-side mistake (e.g. forgetting to fill the encoder
  /// output before calling in).
  #[test]
  fn log_softmax_rejects_empty_buffer_with_nonzero_dims() {
    let err = log_softmax_with_finite_guard(&[], 1, 1).unwrap_err();
    let EmissionsError::Shape(payload) = err else {
      panic!("expected EmissionsError::Shape; got {err:?}");
    };
    assert!(payload.to_string().contains("shape mismatch"));
  }

  /// `raw` longer than the declared `T × V` never panics — the
  /// loop only ever reads the first `T × V` elements — so it
  /// previously returned `Ok` built from a truncated prefix,
  /// silently discarding the trailing elements instead of
  /// signalling the shape mismatch.
  #[test]
  fn log_softmax_rejects_oversized_buffer_instead_of_silently_truncating() {
    let err = log_softmax_with_finite_guard(&[0.0f32, 99.0], 1, 1).unwrap_err();
    let EmissionsError::Shape(payload) = err else {
      panic!("expected EmissionsError::Shape; got {err:?}");
    };
    assert!(
      payload.to_string().contains("shape mismatch"),
      "diagnostic must call out the shape mismatch; got {payload}",
    );
  }

  /// `T = 0` makes the frame loop `0..0`, which never runs
  /// regardless of what `raw` holds — so a non-empty `raw` was
  /// previously discarded whole and `Ok(vec![])` returned,
  /// silently swallowing every element instead of signalling the
  /// shape mismatch.
  #[test]
  fn log_softmax_rejects_zero_t_with_nonempty_buffer_instead_of_silently_discarding() {
    let err = log_softmax_with_finite_guard(&[1.0f32, 2.0, 3.0], 0, 5).unwrap_err();
    let EmissionsError::Shape(payload) = err else {
      panic!("expected EmissionsError::Shape; got {err:?}");
    };
    assert!(payload.to_string().contains("shape mismatch"));
  }

  /// `T = 0` with a correctly-empty `raw` (`0 * V == 0 ==
  /// raw.len()`) is legitimate degenerate input — zero frames in,
  /// zero log-probabilities out — and must stay `Ok`, not be
  /// swept up by the shape guard.
  #[test]
  fn log_softmax_accepts_zero_t_with_empty_buffer() {
    let out = log_softmax_with_finite_guard(&[], 0, 5).expect("ok");
    assert!(out.is_empty());
  }

  /// `T = 0, V = 0` passes the entry-length check trivially (`0 *
  /// 0 == 0 == raw.len()`) and the frame loop `0..0` never runs,
  /// so — before the explicit `V == 0` guard — this silently
  /// returned `Ok(vec![])` instead of rejecting a vocabulary with
  /// no vocab axis at all (not even a blank token). The zero-`T`
  /// case is exactly what let this slip past both the shape check
  /// and the per-row loop.
  #[test]
  fn log_softmax_rejects_zero_vocab_with_zero_t() {
    let err = log_softmax_with_finite_guard(&[], 0, 0).unwrap_err();
    let EmissionsError::Shape(payload) = err else {
      panic!("expected EmissionsError::Shape; got {err:?}");
    };
    assert!(
      payload.to_string().contains("zero-length vocab"),
      "diagnostic must call out the zero-length vocab dim; got {payload}",
    );
  }

  /// `T = 1, V = 0` also passes the entry-length check (`1 * 0 ==
  /// 0 == raw.len()`), then the frame loop's single empty row
  /// (`v == 0` makes every row slice `raw[0..0]`) already fell
  /// through to the shifted-normaliser-non-finite branch
  /// (`sum.ln()` of an empty row's zero-sum is `-inf`) — a typed
  /// `Err`, but a misleading one that names "non-finite
  /// normaliser" instead of the true cause. The explicit `V == 0`
  /// guard now catches this case earlier, with an accurate
  /// message, the same way it closes the silent-`Ok` hole at
  /// `T == 0`.
  #[test]
  fn log_softmax_rejects_zero_vocab_with_positive_t() {
    let err = log_softmax_with_finite_guard(&[], 1, 0).unwrap_err();
    let EmissionsError::Shape(payload) = err else {
      panic!("expected EmissionsError::Shape; got {err:?}");
    };
    assert!(
      payload.to_string().contains("zero-length vocab"),
      "diagnostic must call out the zero-length vocab dim, not the \
 unrelated shifted-normaliser-non-finite path it used to fall \
 through to; got {payload}",
    );
  }

  /// `T * V` overflowing `usize` previously reached
  /// `Vec::with_capacity(t * v)` and `usize::MAX * 2` panicked via
  /// the debug-mode overflow check (release profile would instead
  /// wrap to a small-but-wrong capacity and abort on the resulting
  /// allocation request) — the same overflow discipline
  /// `validate_output_dims` already applies to the ORT-sourced
  /// shape, applied here via `checked_mul` before any arithmetic
  /// on `t`/`v` runs. An unrepresentable product folds into the
  /// shape arm (as it does for `LogProbsTV::new`).
  #[test]
  fn log_softmax_rejects_t_v_product_overflow() {
    let err = log_softmax_with_finite_guard(&[], usize::MAX, 2).unwrap_err();
    let EmissionsError::Shape(payload) = err else {
      panic!("expected EmissionsError::Shape; got {err:?}");
    };
    assert!(
      payload.to_string().contains("shape mismatch"),
      "an unrepresentable T*V product is rejected as a shape error; got {payload}",
    );
  }

  // --- end log_softmax_with_finite_guard dimension validation ---

  // --- ORT output dims validation ---

  #[test]
  fn validate_output_dims_rejects_negative_t() {
    use crate::types::Lang;
    let err = validate_output_dims(-1, 32, 32, &Lang::En).unwrap_err();
    let WorkFailure::Alignment(AlignmentError::ModelInference(payload)) = err else {
      panic!("expected AlignmentFailed");
    };
    let message = payload.message();
    assert!(message.contains("negative time dim"));
  }

  #[test]
  fn validate_output_dims_rejects_zero_v() {
    use crate::types::Lang;
    let err = validate_output_dims(100, 0, 0, &Lang::En).unwrap_err();
    // V=0 is always fatal — model has no vocab axis.
    assert!(matches!(
      err,
      WorkFailure::Alignment(AlignmentError::ModelInference(_))
    ));
  }

  /// A chunk too short to produce any encoder frame must surface
  /// as recoverable `NoAlignmentPath`, not fatal
  /// `ModelInferenceFailed`. The ASR transcript stays alive with
  /// `words: []`.
  #[test]
  fn validate_output_dims_zero_t_with_empty_buffer_is_recoverable_no_alignment_path() {
    use crate::types::Lang;
    let err = validate_output_dims(0, 32, 0, &Lang::En).unwrap_err();
    let WorkFailure::Alignment(AlignmentError::NoAlignmentPath(payload)) = &err else {
      panic!("expected NoAlignmentPath; got {err:?}");
    };
    let message = payload.message();
    assert!(
      message.contains("zero encoder frames"),
      "diagnostic must explain the short-chunk cause; got {message}",
      message = message
    );
  }

  /// T=0 with a non-empty buffer means the model declared zero
  /// frames but returned data anyway — a shape/data
  /// inconsistency that should stay fatal.
  #[test]
  fn validate_output_dims_zero_t_with_nonempty_buffer_stays_fatal() {
    use crate::types::Lang;
    let err = validate_output_dims(0, 32, 5, &Lang::En).unwrap_err();
    let WorkFailure::Alignment(AlignmentError::ModelInference(payload)) = &err else {
      panic!("T=0 with non-empty buffer must stay fatal; got {err:?}");
    };
    let message = payload.message();
    assert!(
      message.contains("shape/data mismatch") || message.contains("buffer has"),
      "diagnostic must call out the shape/data inconsistency; got {message}",
      message = message
    );
  }

  // -------- stride / vocab-dim guards --------

  /// Stride in range — e.g., 16 000-sample chunk at hop=320
  /// gives T=49 (`49 × 320 = 15 680`, 320-sample slack from the
  /// chunk extent). Within the ±2-frame band, accepted.
  #[test]
  fn validate_stride_extent_accepts_typical_under_extent() {
    use crate::types::Lang;
    assert!(validate_stride_extent(49, 320, 16_000, &Lang::En).is_ok());
    // Exact integer match
    assert!(validate_stride_extent(50, 320, 16_000, &Lang::En).is_ok());
    // 1-frame over (within 2-frame slack)
    assert!(validate_stride_extent(51, 320, 16_000, &Lang::En).is_ok());
  }

  /// Stride too small (T overshoots): the model emits more
  /// frames than the chunk could produce, e.g. claimed stride
  /// is 320 but real stride is 160 → T is roughly 2× expected.
  /// Rejected as fatal `ModelInferenceFailed`.
  #[test]
  fn validate_stride_extent_rejects_t_too_large() {
    use crate::types::Lang;
    // 100 frames × 320 = 32 000 sample-equivalents for a
    // 16 000-sample chunk. Way past the upper bound (16 640).
    let err = validate_stride_extent(100, 320, 16_000, &Lang::En).unwrap_err();
    let WorkFailure::Alignment(AlignmentError::ModelInference(payload)) = err else {
      panic!("expected AlignmentFailed");
    };
    let message = payload.message();
    assert!(
      message.contains("smaller stride"),
      "diagnostic must call out the smaller-stride case; got {message}",
      message = message
    );
  }

  /// Stride too large (T undershoots): the model emits far
  /// fewer frames than the input audio supports, e.g. claimed
  /// stride is 320 but real stride is 640. Without this check,
  /// `compose_words` would compress every word into the first
  /// half of the chunk's audio. Rejected as fatal
  /// `ModelInferenceFailed`.
  #[test]
  fn validate_stride_extent_rejects_t_too_small() {
    use crate::types::Lang;
    // 25 frames × 320 = 8 000 sample-equivalents for a 16 000-
    // sample chunk. Half the expected — far below the lower
    // bound (15 360 = 16 000 − 640).
    let err = validate_stride_extent(25, 320, 16_000, &Lang::En).unwrap_err();
    let WorkFailure::Alignment(AlignmentError::ModelInference(payload)) = err else {
      panic!("expected AlignmentFailed");
    };
    let message = payload.message();
    assert!(
      message.contains("larger stride"),
      "diagnostic must call out the larger-stride case; got {message}",
      message = message
    );
  }

  /// Very short chunks where the slack is comparable to the
  /// chunk extent — the lower bound saturates to 0 and small
  /// `T` values pass. (T=0 itself is routed to recoverable
  /// `NoAlignmentPath` upstream by `validate_output_dims`.)
  #[test]
  fn validate_stride_extent_accepts_very_short_chunk_with_small_t() {
    use crate::types::Lang;
    // 200-sample chunk, hop=320 → slack=640, lower=0.
    // T=1 → frame_extent=320, within [0, 840]. Accepted.
    assert!(validate_stride_extent(1, 320, 200, &Lang::En).is_ok());
  }

  /// Vocab-dim equality: model output V matches tokenizer
  /// vocab size → accepted.
  #[test]
  fn validate_vocab_dim_accepts_exact_match() {
    use crate::types::Lang;
    assert!(validate_vocab_dim(32, 32, &Lang::En).is_ok());
  }

  /// Vocab-dim mismatch: model output V is larger than the
  /// tokenizer's vocab. Rejected as fatal — Viterbi would
  /// otherwise read posteriors from columns the tokenizer
  /// thinks correspond to the wrong tokens.
  #[test]
  fn validate_vocab_dim_rejects_oversized_model_output() {
    use crate::types::Lang;
    let err = validate_vocab_dim(1024, 32, &Lang::En).unwrap_err();
    let WorkFailure::Alignment(AlignmentError::ModelInference(payload)) = err else {
      panic!("expected AlignmentFailed");
    };
    let message = payload.message();
    assert!(
      message.contains("doesn't match tokenizer vocab"),
      "diagnostic must call out the vocab mismatch; got {message}",
      message = message
    );
  }

  /// Vocab-dim mismatch: model output V is smaller than the
  /// tokenizer's vocab. Same rejection.
  #[test]
  fn validate_vocab_dim_rejects_undersized_model_output() {
    use crate::types::Lang;
    let err = validate_vocab_dim(16, 32, &Lang::En).unwrap_err();
    assert!(matches!(
      err,
      WorkFailure::Alignment(AlignmentError::ModelInference(_))
    ));
  }

  // -------- end stride / vocab-dim guards --------

  #[test]
  fn validate_output_dims_rejects_buffer_length_mismatch() {
    use crate::types::Lang;
    // Declared T=10, V=4 → 40 elements; provided buffer = 39.
    let err = validate_output_dims(10, 4, 39, &Lang::En).unwrap_err();
    let WorkFailure::Alignment(payload) = err else {
      panic!("expected AlignmentFailed");
    };
    let message = payload.to_string();
    assert!(
      message.contains("doesn't match"),
      "must call out length mismatch; got {message}",
      message = message
    );
  }

  /// 32-bit-only path: usize is 32 bits, so `i64` of 1 << 33
  /// can't fit. We test the overflow logic indirectly by
  /// asking for `T * V` larger than usize::MAX; on aarch64
  /// (64-bit) we'd need an astronomical product, so this
  /// targets the `checked_mul` branch with two values whose
  /// product overflows. usize::MAX ≈ 1.8e19 on 64-bit; we use
  /// √max + 1 each.
  #[test]
  fn validate_output_dims_rejects_t_v_product_overflow() {
    use crate::types::Lang;
    // Two large values whose product overflows usize on any
    // platform. 2^32 × 2^32 = 2^64 > usize::MAX on 64-bit
    // (overflow); same on 32-bit (overflow much earlier).
    let big = i64::from(u32::MAX) + 1; // 2^32
    let err = validate_output_dims(big, big, 0, &Lang::En).unwrap_err();
    let WorkFailure::Alignment(payload) = err else {
      panic!("expected AlignmentFailed");
    };
    let message = payload.to_string();
    assert!(
      message.contains("overflow") || message.contains("doesn't fit"),
      "must call out overflow; got {message}",
      message = message
    );
  }

  #[test]
  fn validate_output_dims_accepts_well_formed_shape() {
    use crate::types::Lang;
    let (t, v) = validate_output_dims(1500, 32, 1500 * 32, &Lang::En).expect("ok");
    assert_eq!(t, 1500);
    assert_eq!(v, 32);
  }

  /// Multi-frame: a NaN in frame 2 surfaces with `frame 2` in
  /// the message. Locks in that the frame index is precise
  /// (helpful for debugging upstream backend issues).
  #[test]
  fn log_softmax_locates_nan_to_specific_frame() {
    // 3 frames × 2 vocab; frame 2's first element is NaN.
    let raw = vec![0.0_f32, 0.1, 0.0, 0.1, f32::NAN, 0.1];
    let err = log_softmax_with_finite_guard(&raw, 3, 2).unwrap_err();
    let EmissionsError::Numeric(payload) = err else {
      panic!("expected EmissionsError::Numeric; got {err:?}");
    };
    let message = payload.to_string();
    assert!(
      message.contains("frame 2"),
      "must locate the bad frame; got {message}",
      message = message
    );
  }

  // Note: the centring / scale and empty-input behaviour tests
  // moved to `super::normalize::tests` after normalisation was
  // pulled up the call stack into `Aligner::align`.
  // `encode_log_softmax` no longer normalises, so its tests
  // here cover only the reductions and the input-validation
  // boundary it does still own.
}