franken_ocr 0.8.0

//! SAM ViT-B encoder forward ([SPEC-040..046], PROPOSED_ARCHITECTURE.md §6.3).
//! Vision config provenance: SAM ViT-B width 768 / depth 12 / 12 heads /
//! global attention at [2,5,8,11], image size 1024 ([SPEC-017]).
//!
//! Real fp32 forward. Patch-embed Conv2d k16s16 -> 64x64 tokens (width 768);
//! learned `pos_embed` (1,64,64,768) bicubic-interpolated to the runtime grid;
//! 12 [`Block`]s with window attention (`window=14`, OQ-15) except global at
//! `[2,5,8,11]`; decomposed relative-position bias added to the SDPA logits; the
//! neck (Conv2d 768->256 k1 -> LayerNorm2d -> Conv2d 256->256 k3 p1 ->
//! LayerNorm2d) then two stride-2 downsamples (`net_2` 256->512, `net_3`
//! 512->1024) -> `[B, 1024, 16, 16]`, returned as a row-major `Mat` of shape
//! `[1024, 16*16]` (channels x flattened spatial), matching the `flatten(2)`
//! layout the bridge concatenates (OQ-6: `local_features_1.flatten(2)`).
//!
//! Weights are owned by the parallel weights wave; this module operates over a
//! [`SamWeights`] parameter bundle so the full math is unit-testable on tiny
//! synthetic inputs with no model present. The public [`forward`] entrypoint
//! adapts a loaded [`Weights`] into a [`SamWeights`] through the `.focrq`
//! reader's named-tensor accessors (bd-1es.3) and runs the real SAM forward.

use rayon::prelude::*;

use super::nn;
use super::tensor::Mat;
use super::weights::Weights;
use crate::error::{FocrError, FocrResult};

// Clock seam: `std::time::Instant` traps on wasm32-unknown-unknown; `web-time`
// re-exports std's types on native targets, so native behavior is unchanged.
#[cfg(not(target_arch = "wasm32"))]
use std::time::Instant;
#[cfg(target_arch = "wasm32")]
use web_time::Instant;

// ── fixed SAM-ViT-B geometry ([SPEC-040..046]) ─────────────────────────────

/// Patch embedding / transformer width.
pub const EMBED_DIM: usize = 768;
/// Number of transformer blocks.
pub const DEPTH: usize = 12;
/// Attention heads per block.
pub const NUM_HEADS: usize = 12;
/// Per-head channel count (`768 / 12 = 64`).
pub const HEAD_DIM: usize = EMBED_DIM / NUM_HEADS;
/// Patch-embed kernel / stride (`16`), so `1024 -> 64`.
pub const PATCH: usize = 16;
/// Window size for non-global blocks (OQ-15: 14).
pub const WINDOW: usize = 14;
/// Block indices that run global (full-grid) attention.
pub const GLOBAL_BLOCKS: [usize; 4] = [2, 5, 8, 11];
/// Neck / `out_chans` channel count.
pub const NECK_CH: usize = 256;
/// `net_2` output channels.
pub const NET2_CH: usize = 512;
/// `net_3` output channels (the returned feature width).
pub const OUT_CH: usize = 1024;
/// LayerNorm eps for the transformer norms and `LayerNorm2d`.
pub const LN_EPS: f32 = 1e-6;
/// MLP hidden = `dim * mlp_ratio` (`mlp_ratio = 4`).
pub const MLP_HIDDEN: usize = EMBED_DIM * 4;

fn checked_shape_mul(context: &str, lhs: usize, rhs: usize, expression: &str) -> FocrResult<usize> {
    lhs.checked_mul(rhs).ok_or_else(|| {
        FocrError::Other(anyhow::anyhow!(
            "{context}: usize overflow computing {expression} ({lhs} * {rhs})"
        ))
    })
}

fn checked_shape_add(context: &str, lhs: usize, rhs: usize, expression: &str) -> FocrResult<usize> {
    lhs.checked_add(rhs).ok_or_else(|| {
        FocrError::Other(anyhow::anyhow!(
            "{context}: usize overflow computing {expression} ({lhs} + {rhs})"
        ))
    })
}

fn checked_shape_sub(context: &str, lhs: usize, rhs: usize, expression: &str) -> FocrResult<usize> {
    lhs.checked_sub(rhs).ok_or_else(|| {
        FocrError::Other(anyhow::anyhow!(
            "{context}: usize underflow computing {expression} ({lhs} - {rhs})"
        ))
    })
}

fn checked_nchw_len(
    context: &str,
    ch: usize,
    h: usize,
    w: usize,
    expression: &str,
) -> FocrResult<usize> {
    let hw = checked_shape_mul(context, h, w, "h*w")?;
    checked_shape_mul(context, ch, hw, expression)
}

fn checked_conv_weight_len(
    context: &str,
    out_ch: usize,
    in_ch: usize,
    kh: usize,
    kw: usize,
) -> FocrResult<usize> {
    let out_in = checked_shape_mul(context, out_ch, in_ch, "out_ch*in_ch")?;
    let kernel = checked_shape_mul(context, kh, kw, "kh*kw")?;
    checked_shape_mul(context, out_in, kernel, "out_ch*in_ch*kh*kw")
}

// ── parameter bundles ──────────────────────────────────────────────────────

/// A `nn.Linear` parameter pair, weight pre-transposed to the GEMM-ready
/// `[in, out]` row-major layout plus a length-`out` bias.
///
/// PyTorch stores `Linear.weight` as `[out_features, in_features]`;
/// [`Self::from_row_major`] transposes it ONCE at construction so `apply` is a
/// straight matmul — bd-av64.10 measured the old transpose-at-apply-time as
/// hundreds of MB of pure data movement per vision forward, and TrOMR's AR
/// loop paid it PER DECODE STEP. Same floats in a different order ⇒ outputs
/// byte-identical. The weight field is private so a mis-shaped Linear is
/// unrepresentable outside the validating constructor.
#[derive(Debug, Clone)]
pub struct Linear {
    /// Weight pre-transposed to `[in_, out]` row-major.
    wt: Mat,
    /// Length-`out` bias (may be empty for `bias=False`).
    pub b: Vec<f32>,
    /// Output features.
    pub out: usize,
    /// Input features.
    pub in_: usize,
}

impl Linear {
    /// Build from a PyTorch `[out, in]` row-major weight, transposing once and
    /// validating every length at construction (fail-fast, so `apply` never
    /// discovers a malformed bundle mid-forward).
    ///
    /// # Errors
    /// [`FocrError::Other`] when `w.len() != out*in` or a non-empty bias has
    /// `b.len() != out`.
    pub fn from_row_major(w: &[f32], b: Vec<f32>, out: usize, in_: usize) -> FocrResult<Self> {
        let expected_weight_len = checked_shape_mul("vision_sam linear", out, in_, "out*in")?;
        if w.len() != expected_weight_len {
            return Err(FocrError::Other(anyhow::anyhow!(
                "vision_sam linear: weight len {} != out*in {}",
                w.len(),
                expected_weight_len
            )));
        }
        if !b.is_empty() && b.len() != out {
            return Err(FocrError::Other(anyhow::anyhow!(
                "vision_sam linear: bias len {} != out_features {}",
                b.len(),
                out
            )));
        }
        // w is [out, in]; transpose to [in, out] so matmul([m,in],[in,out]).
        let wt = Mat::from_vec(in_, out, transpose(w, out, in_));
        Ok(Self { wt, b, out, in_ })
    }

    /// `y[m,out] = x[m,in] @ w^T + b`. `x.cols` must equal `self.in_`. Also the
    /// GOT `mm_projector_vary` connector (a plain Linear(1024→1024)+bias).
    ///
    /// # Errors
    /// [`FocrError::Other`] if `x.cols != self.in_` or the public shape
    /// metadata was corrupted after construction.
    pub fn apply(&self, x: &Mat) -> FocrResult<Mat> {
        if x.cols != self.in_ {
            return Err(FocrError::Other(anyhow::anyhow!(
                "vision_sam linear: input cols {} != expected in_features {}",
                x.cols,
                self.in_
            )));
        }
        if self.wt.rows != self.in_ || self.wt.cols != self.out {
            return Err(FocrError::Other(anyhow::anyhow!(
                "vision_sam linear: pretransposed weight shape [{}, {}] != [in,out] [{}, {}]",
                self.wt.rows,
                self.wt.cols,
                self.in_,
                self.out
            )));
        }
        if !self.b.is_empty() && self.b.len() != self.out {
            return Err(FocrError::Other(anyhow::anyhow!(
                "vision_sam linear: bias len {} != out_features {}",
                self.b.len(),
                self.out
            )));
        }
        let mut y = nn::matmul(x, &self.wt)?;
        if !self.b.is_empty() {
            for r in 0..y.rows {
                let row = y.row_mut(r);
                for (c, v) in row.iter_mut().enumerate() {
                    *v += self.b[c];
                }
            }
        }
        Ok(y)
    }
}

/// A LayerNorm affine pair over `cols` features.
#[derive(Debug, Clone)]
pub struct LayerNormP {
    /// Length-`cols` gain.
    pub w: Vec<f32>,
    /// Length-`cols` shift.
    pub b: Vec<f32>,
}

/// A 2-D convolution weight (`[out_ch, in_ch, kh, kw]` row-major) + optional
/// bias.
#[derive(Debug, Clone)]
pub struct Conv {
    /// `[out_ch, in_ch, kh, kw]` row-major.
    pub w: Vec<f32>,
    /// Optional length-`out_ch` bias.
    pub b: Option<Vec<f32>>,
    /// Output channels.
    pub out_ch: usize,
    /// Input channels.
    pub in_ch: usize,
    /// Kernel height.
    pub kh: usize,
    /// Kernel width.
    pub kw: usize,
}

/// Per-block attention parameters (qkv fused, output proj, rel-pos tables).
#[derive(Debug, Clone)]
pub struct AttnP {
    /// Fused qkv linear: `[3*dim, dim]` weight, `3*dim` bias.
    pub qkv: Linear,
    /// Output projection `[dim, dim]`.
    pub proj: Linear,
    /// `rel_pos_h`: `[2*size_h - 1, head_dim]` row-major.
    pub rel_pos_h: Vec<f32>,
    /// `rel_pos_w`: `[2*size_w - 1, head_dim]` row-major.
    pub rel_pos_w: Vec<f32>,
    /// Rel-pos table spatial size along H (window for windowed blocks, 64 for
    /// global). `rel_pos_h` has `2*size_h - 1` rows.
    pub size_h: usize,
    /// Rel-pos table spatial size along W.
    pub size_w: usize,
}

/// One transformer block's parameters.
#[derive(Debug, Clone)]
pub struct BlockP {
    /// `norm1` (pre-attention).
    pub norm1: LayerNormP,
    /// Attention.
    pub attn: AttnP,
    /// `norm2` (pre-MLP).
    pub norm2: LayerNormP,
    /// MLP `lin1` (`dim -> 4*dim`).
    pub lin1: Linear,
    /// MLP `lin2` (`4*dim -> dim`).
    pub lin2: Linear,
    /// Effective window size (0 => global block).
    pub window: usize,
}

/// The full SAM-ViT-B parameter set.
#[derive(Debug, Clone)]
pub struct SamWeights {
    /// Patch-embed conv (`3 -> 768`, k16 s16).
    pub patch_embed: Conv,
    /// Learned abs pos-embed, row-major `[grid_h, grid_w, dim]` (canonical
    /// `64x64x768`).
    pub pos_embed: Vec<f32>,
    /// Pos-embed source grid height (canonical 64).
    pub pos_grid_h: usize,
    /// Pos-embed source grid width (canonical 64).
    pub pos_grid_w: usize,
    /// The 12 transformer blocks.
    pub blocks: Vec<BlockP>,
    /// Neck conv1 (`768 -> 256`, k1, no bias).
    pub neck_conv1: Conv,
    /// Neck LayerNorm2d #1 (`256`).
    pub neck_ln1: LayerNormP,
    /// Neck conv2 (`256 -> 256`, k3 p1, no bias).
    pub neck_conv2: Conv,
    /// Neck LayerNorm2d #2 (`256`).
    pub neck_ln2: LayerNormP,
    /// `net_2` (`256 -> 512`, k3 s2 p1, no bias).
    pub net2: Conv,
    /// `net_3` (`512 -> 1024`, k3 s2 p1, no bias).
    pub net3: Conv,
}

// ── public entrypoints ─────────────────────────────────────────────────────

/// Run the SAM tower over a normalized `[3, H, W]` image, returning the `x3`
/// 1024-channel feature flattened to `[1024, (H/64spatial)^2]` — i.e. the
/// `net_3` output `[1, 1024, 16, 16]` reshaped channel-major
/// (`flatten(2)` layout, OQ-6).
///
/// # Errors
/// Propagates accessor errors from building the [`SamWeights`] (e.g. a missing
/// or mis-shaped `model.sam_model.*` tensor), the input-shape checks above, and
/// whatever [`forward_with`] returns. The real math lives in [`forward_with`],
/// which is exercised by the unit tests below.
pub fn forward(weights: &Weights, image: &Mat) -> FocrResult<Mat> {
    forward_prefix(weights, image, "model.sam_model")
}

/// [`forward`] with an explicit `.focrq` tensor-name prefix for the SAM-family
/// vision tower — Baidu Unlimited-OCR uses `model.sam_model` (the default), GOT-OCR2
/// uses `model.vision_tower_high`; the leaf tensor names + all geometry are identical.
///
/// # Errors
/// As [`forward`] — the first vision-stage error (missing/mis-shaped tensor or a
/// kernel failure).
pub fn forward_prefix(weights: &Weights, image: &Mat, prefix: &str) -> FocrResult<Mat> {
    if image.rows != 3 {
        return Err(FocrError::Other(anyhow::anyhow!(
            "vision_sam::forward: expected 3 input channels, got {}",
            image.rows
        )));
    }
    // Base-mode vision input is square: image.cols == H*W with H == W.
    let side = (image.cols as f64).sqrt() as usize;
    if side * side != image.cols {
        return Err(FocrError::Other(anyhow::anyhow!(
            "vision_sam::forward: image.cols {} is not a perfect square",
            image.cols
        )));
    }
    let th = Instant::now();
    let w = sam_weights_from(weights, prefix)?;
    super::timing_log(&format!(
        "    sam.hydrate {:.2}s",
        th.elapsed().as_secs_f64()
    ));
    let tf = Instant::now();
    let out = forward_with(&w, image, side, side);
    super::timing_log(&format!(
        "    sam.forward {:.2}s",
        tf.elapsed().as_secs_f64()
    ));
    out
}

/// Build a [`SamWeights`] from the named `{prefix}.*` tensors in `weights`
/// (BF16→f32 widened at the accessor). Dims are read from each tensor's shape;
/// rel-pos table sizes are derived from the stored rows. `pub(crate)` so the
/// batch spine (bd-1azu.10) hydrates ONCE per batch instead of once per view.
pub(crate) fn sam_weights_from(weights: &Weights, prefix: &str) -> FocrResult<SamWeights> {
    // The embed/neck/per-block builders are the SAME functions the streamed
    // forward uses (bd-4l71), so the two hydration modes cannot drift; only
    // WHERE the hydrated blocks live differs. Hydration ORDER (patch/pos ->
    // blocks -> neck) is the historical one, so error precedence on malformed
    // artifacts is unchanged.
    let (patch_embed, pos_embed, pos_grid_h, pos_grid_w) = sam_embed_from(weights, prefix)?;
    let mut blocks = Vec::with_capacity(DEPTH);
    for i in 0..DEPTH {
        blocks.push(sam_block_from(weights, prefix, i)?);
    }
    let mut w = sam_neck_shell(
        weights,
        prefix,
        patch_embed,
        pos_embed,
        pos_grid_h,
        pos_grid_w,
    )?;
    w.blocks = blocks;
    Ok(w)
}

/// Hydrate ONE SAM transformer block from the `{prefix}.blocks.{i}.*` tensors —
/// the exact per-block build [`sam_weights_from`] performs (it calls this), so
/// the streamed per-block forward (bd-4l71 wasm-lane residency) and the cached
/// whole-tower hydration CANNOT drift.
pub(crate) fn sam_block_from(weights: &Weights, prefix: &str, i: usize) -> FocrResult<BlockP> {
    let flat = |n: &str| weights.vec(n);
    let ln = |n: &str| -> FocrResult<LayerNormP> {
        Ok(LayerNormP {
            w: flat(&format!("{n}.weight"))?,
            b: flat(&format!("{n}.bias"))?,
        })
    };
    let b = format!("{prefix}.blocks.{i}");
    let window = if GLOBAL_BLOCKS.contains(&i) {
        0
    } else {
        WINDOW
    };
    let qkv_name = format!("{b}.attn.qkv.weight");
    let proj_name = format!("{b}.attn.proj.weight");
    let rph_name = format!("{b}.attn.rel_pos_h");
    let rpw_name = format!("{b}.attn.rel_pos_w");
    let lin1_name = format!("{b}.mlp.lin1.weight");
    let lin2_name = format!("{b}.mlp.lin2.weight");
    let (q_out, q_in) = tensor_rank2_shape(weights, &qkv_name)?;
    let (proj_out, proj_in) = tensor_rank2_shape(weights, &proj_name)?;
    let (rph_rows, _rph_cols) = tensor_rank2_shape(weights, &rph_name)?;
    let (rpw_rows, _rpw_cols) = tensor_rank2_shape(weights, &rpw_name)?;
    let (lin1_out, lin1_in) = tensor_rank2_shape(weights, &lin1_name)?;
    let (lin2_out, lin2_in) = tensor_rank2_shape(weights, &lin2_name)?;
    Ok(BlockP {
        norm1: ln(&format!("{b}.norm1"))?,
        attn: AttnP {
            qkv: Linear::from_row_major(
                &flat(&qkv_name)?,
                flat(&format!("{b}.attn.qkv.bias"))?,
                q_out,
                q_in,
            )?,
            proj: Linear::from_row_major(
                &flat(&proj_name)?,
                flat(&format!("{b}.attn.proj.bias"))?,
                proj_out,
                proj_in,
            )?,
            rel_pos_h: flat(&rph_name)?,
            rel_pos_w: flat(&rpw_name)?,
            size_h: rph_rows.div_ceil(2),
            size_w: rpw_rows.div_ceil(2),
        },
        norm2: ln(&format!("{b}.norm2"))?,
        lin1: Linear::from_row_major(
            &flat(&lin1_name)?,
            flat(&format!("{b}.mlp.lin1.bias"))?,
            lin1_out,
            lin1_in,
        )?,
        lin2: Linear::from_row_major(
            &flat(&lin2_name)?,
            flat(&format!("{b}.mlp.lin2.bias"))?,
            lin2_out,
            lin2_in,
        )?,
        window,
    })
}

/// Hydrate the patch-embed conv + abs pos-embed (the pre-block head parts, in
/// the historical order). Shared by [`sam_weights_from`] and [`sam_head_from`].
fn sam_embed_from(weights: &Weights, prefix: &str) -> FocrResult<(Conv, Vec<f32>, usize, usize)> {
    let p = prefix;
    let flat = |n: &str| weights.vec(n);
    let pe = format!("{p}.patch_embed.proj.weight");
    let (pe_out, pe_in, pe_kh, pe_kw) = tensor_rank4_shape(weights, &pe)?;
    let patch_embed = Conv {
        w: flat(&pe)?,
        b: Some(flat(&format!("{p}.patch_embed.proj.bias"))?),
        out_ch: pe_out,
        in_ch: pe_in,
        kh: pe_kh,
        kw: pe_kw,
    };
    let pos_name = format!("{p}.pos_embed");
    let (pgh, pgw) = tensor_pos_grid_shape(weights, &pos_name)?;
    let pos_embed = flat(&pos_name)?;
    Ok((patch_embed, pos_embed, pgh, pgw))
}

/// Hydrate the neck/net convs and assemble a [`SamWeights`] with an EMPTY
/// `blocks` vec (the caller supplies blocks, or streams them per use).
fn sam_neck_shell(
    weights: &Weights,
    prefix: &str,
    patch_embed: Conv,
    pos_embed: Vec<f32>,
    pos_grid_h: usize,
    pos_grid_w: usize,
) -> FocrResult<SamWeights> {
    let p = prefix;
    let flat = |n: &str| weights.vec(n);
    let conv = |n: &str, bias: bool| -> FocrResult<Conv> {
        let d = tensor_min_rank_shape(weights, n, 2)?;
        Ok(Conv {
            w: flat(n)?,
            b: if bias {
                Some(flat(&n.replace(".weight", ".bias"))?)
            } else {
                None
            },
            out_ch: d[0],
            in_ch: d[1],
            kh: d.get(2).copied().unwrap_or(1),
            kw: d.get(3).copied().unwrap_or(1),
        })
    };
    let ln = |n: &str| -> FocrResult<LayerNormP> {
        Ok(LayerNormP {
            w: flat(&format!("{n}.weight"))?,
            b: flat(&format!("{n}.bias"))?,
        })
    };
    Ok(SamWeights {
        patch_embed,
        pos_embed,
        pos_grid_h,
        pos_grid_w,
        blocks: Vec::new(),
        neck_conv1: conv(&format!("{p}.neck.0.weight"), false)?,
        neck_ln1: ln(&format!("{p}.neck.1"))?,
        neck_conv2: conv(&format!("{p}.neck.2.weight"), false)?,
        neck_ln2: ln(&format!("{p}.neck.3"))?,
        net2: conv(&format!("{p}.net_2.weight"), false)?,
        net3: conv(&format!("{p}.net_3.weight"), false)?,
    })
}

/// Hydrate ONLY the small non-block SAM parameters (patch embed, pos embed,
/// neck/net convs — ~10 MB f32) with an EMPTY `blocks` vec: the head of the
/// streamed per-block forward (bd-4l71 wasm-lane residency).
pub(crate) fn sam_head_from(weights: &Weights, prefix: &str) -> FocrResult<SamWeights> {
    let (patch_embed, pos_embed, pgh, pgw) = sam_embed_from(weights, prefix)?;
    sam_neck_shell(weights, prefix, patch_embed, pos_embed, pgh, pgw)
}

fn tensor_min_rank_shape(weights: &Weights, name: &str, min_rank: usize) -> FocrResult<Vec<usize>> {
    let view = weights.tensor(name)?;
    if view.shape.len() < min_rank {
        return Err(FocrError::FormatMismatch(format!(
            "tensor {name:?} has rank {}; expected at least {min_rank}",
            view.shape.len()
        )));
    }
    Ok(view.shape.to_vec())
}

fn tensor_rank2_shape(weights: &Weights, name: &str) -> FocrResult<(usize, usize)> {
    let view = weights.tensor(name)?;
    let [rows, cols] = view.shape else {
        return Err(FocrError::FormatMismatch(format!(
            "tensor {name:?} has rank {}; expected 2 ([rows, cols])",
            view.shape.len()
        )));
    };
    Ok((*rows, *cols))
}

fn tensor_rank4_shape(weights: &Weights, name: &str) -> FocrResult<(usize, usize, usize, usize)> {
    let view = weights.tensor(name)?;
    let [out_ch, in_ch, kh, kw] = view.shape else {
        return Err(FocrError::FormatMismatch(format!(
            "tensor {name:?} has rank {}; expected 4 ([out_ch, in_ch, kh, kw])",
            view.shape.len()
        )));
    };
    Ok((*out_ch, *in_ch, *kh, *kw))
}

fn tensor_pos_grid_shape(weights: &Weights, name: &str) -> FocrResult<(usize, usize)> {
    let view = weights.tensor(name)?;
    match view.shape {
        [_h0, h, w, _c] => Ok((*h, *w)),
        [h, w, _c] => Ok((*h, *w)),
        shape => Err(FocrError::FormatMismatch(format!(
            "tensor {name:?} has rank {}; expected 3 ([h, w, c]) or 4 ([1, h, w, c]), got {shape:?}",
            shape.len()
        ))),
    }
}

/// Run the SAM tower with an explicit parameter bundle over a `[3, H, W]`
/// image (laid out `[in_ch=3, H, W]` row-major in `image.data`, `image.rows=3`,
/// `image.cols=H*W` with `H == W` a multiple of `PATCH`).
///
/// Returns the `net_3` feature as a row-major `[OUT_CH, gh3*gw3]` [`Mat`]
/// (channel-major / `flatten(2)` order).
///
/// # Errors
/// [`FocrError::Other`] on a shape contract violation (non-3-channel input,
/// non-square or non-`PATCH`-divisible spatial dims, or a kernel rejection).
pub fn forward_with(w: &SamWeights, image: &Mat, h: usize, win: usize) -> FocrResult<Mat> {
    forward_core(
        w,
        image,
        h,
        win,
        w.blocks.len(),
        &mut |i| Ok(std::borrow::Cow::Borrowed(&w.blocks[i])),
        false,
    )
}

/// [`forward`] semantics WITHOUT retaining the whole ~0.4 GB f32 SAM tower
/// (bd-4l71 wasm-lane residency): the small head hydrates once
/// ([`sam_head_from`]), then each block is hydrated ([`sam_block_from`] — the
/// SAME builder the cached path uses), run, and DROPPED. Global-attention
/// blocks run the bounded-scratch query-slab kernel
/// ([`attention_global_bounded`], bit-identical — see its gate) so the
/// transient working set stays tens of MB instead of ~0.8 GB.
///
/// # Errors
/// As [`forward_prefix`].
pub(crate) fn forward_streamed(weights: &Weights, image: &Mat, prefix: &str) -> FocrResult<Mat> {
    if image.rows != 3 {
        return Err(FocrError::Other(anyhow::anyhow!(
            "vision_sam::forward: expected 3 input channels, got {}",
            image.rows
        )));
    }
    let side = (image.cols as f64).sqrt() as usize;
    if side * side != image.cols {
        return Err(FocrError::Other(anyhow::anyhow!(
            "vision_sam::forward: image.cols {} is not a perfect square",
            image.cols
        )));
    }
    let th = Instant::now();
    let head = sam_head_from(weights, prefix)?;
    super::timing_log(&format!(
        "    sam.hydrate(head) {:.2}s",
        th.elapsed().as_secs_f64()
    ));
    let tf = Instant::now();
    let out = forward_core(
        &head,
        image,
        side,
        side,
        DEPTH,
        &mut |i| Ok(std::borrow::Cow::Owned(sam_block_from(weights, prefix, i)?)),
        true,
    );
    super::timing_log(&format!(
        "    sam.forward(streamed) {:.2}s",
        tf.elapsed().as_secs_f64()
    ));
    out
}

/// [`forward_streamed`] over several views at once, hydrating each block ONCE
/// and running EVERY view through it before dropping it (bd-K2 per-view
/// hydration hoist). The head hydrates once as well.
///
/// The `i`-th result is bit-identical to `forward_streamed(weights, images[i],
/// prefix)`: only the loop nest changes (views-inner instead of views-outer),
/// and the views are mathematically independent — see [`forward_core_views`]
/// and the `streamed_views_is_bit_identical_to_per_view` gate.
///
/// Cost: one live activation per view (`[gh·gw, 768]` f32) instead of one, in
/// exchange for `V`× less bf16→f32 block-hydration traffic.
///
/// # Errors
/// As [`forward_streamed`], plus an empty view list.
pub(crate) fn forward_streamed_views(
    weights: &Weights,
    images: &[&Mat],
    prefix: &str,
) -> FocrResult<Vec<Mat>> {
    if images.is_empty() {
        return Err(FocrError::Other(anyhow::anyhow!(
            "vision_sam::forward_streamed_views: empty view list"
        )));
    }
    let mut dims = Vec::with_capacity(images.len());
    for image in images {
        if image.rows != 3 {
            return Err(FocrError::Other(anyhow::anyhow!(
                "vision_sam::forward: expected 3 input channels, got {}",
                image.rows
            )));
        }
        let side = (image.cols as f64).sqrt() as usize;
        if side * side != image.cols {
            return Err(FocrError::Other(anyhow::anyhow!(
                "vision_sam::forward: image.cols {} is not a perfect square",
                image.cols
            )));
        }
        dims.push((side, side));
    }
    let th = Instant::now();
    let head = sam_head_from(weights, prefix)?;
    super::timing_log(&format!(
        "    sam.hydrate(head) {:.2}s",
        th.elapsed().as_secs_f64()
    ));
    let tf = Instant::now();
    let out = forward_core_views(
        &head,
        images,
        &dims,
        DEPTH,
        &mut |i| Ok(std::borrow::Cow::Owned(sam_block_from(weights, prefix, i)?)),
        true,
    );
    super::timing_log(&format!(
        "    sam.forward(streamed, {} views) {:.2}s",
        images.len(),
        tf.elapsed().as_secs_f64()
    ));
    out
}

/// The shared SAM tower core: patch embed + pos embed, `depth` transformer
/// blocks obtained from `block_at` (borrowed from a cached [`SamWeights`] or
/// hydrated per use and dropped), then the neck/net convs. ONE body serves the
/// cached and streamed forwards, so they cannot drift; `low_mem` selects the
/// bounded-scratch global-attention kernel (bit-identical, gated).
#[allow(clippy::too_many_arguments)]
fn forward_core<'w>(
    w: &SamWeights,
    image: &Mat,
    h: usize,
    win: usize,
    depth: usize,
    block_at: &mut dyn FnMut(usize) -> FocrResult<std::borrow::Cow<'w, BlockP>>,
    low_mem: bool,
) -> FocrResult<Mat> {
    let mut out = forward_core_views(
        w,
        std::slice::from_ref(&image),
        &[(h, win)],
        depth,
        block_at,
        low_mem,
    )?;
    Ok(out.remove(0))
}

/// [`forward_core`] over `V` independent views with the **view loop inside the
/// block loop** (bd-K2 per-view hydration hoist): each block is obtained ONCE
/// from `block_at` and applied to every view's activation before the next block
/// is requested. For the streamed lane that turns `V` full hydration sweeps of
/// the bf16 blob into ONE.
///
/// Bit-identity is by construction, not by measurement: view `i`'s activation
/// only ever meets view `i`'s own data, and it passes through exactly the same
/// sequence of kernel calls with exactly the same shapes as it would in the
/// views-outer nest ([`forward_core`] is literally this function at `V == 1`).
/// Reordering *independent* computations cannot move a bit; nothing here is
/// reassociated, restacked, or batched.
///
/// Views may differ in geometry (`dims[i]` is that view's `(H, W)`), since every
/// stage runs per view with that view's own grid.
fn forward_core_views<'w>(
    w: &SamWeights,
    images: &[&Mat],
    dims: &[(usize, usize)],
    depth: usize,
    block_at: &mut dyn FnMut(usize) -> FocrResult<std::borrow::Cow<'w, BlockP>>,
    low_mem: bool,
) -> FocrResult<Vec<Mat>> {
    if images.len() != dims.len() {
        return Err(FocrError::Other(anyhow::anyhow!(
            "vision_sam: {} views but {} geometries",
            images.len(),
            dims.len()
        )));
    }
    let dim = w.patch_embed.out_ch;
    // ── per-view validation + patch embed + abs pos-embed. ──────────────────
    // (Identical checks, in the identical order, to the single-view core.)
    let mut grids = Vec::with_capacity(images.len());
    let mut xs = Vec::with_capacity(images.len());
    for (image, &(h, win)) in images.iter().zip(dims.iter()) {
        if image.rows != 3 {
            return Err(FocrError::Other(anyhow::anyhow!(
                "vision_sam: expected 3 input channels, got {}",
                image.rows
            )));
        }
        if h == 0 || win == 0 {
            return Err(FocrError::Other(anyhow::anyhow!(
                "vision_sam: spatial dims ({h},{win}) must be non-zero"
            )));
        }
        let expected_cols = checked_shape_mul("vision_sam", h, win, "H*W")?;
        if image.cols != expected_cols {
            return Err(FocrError::Other(anyhow::anyhow!(
                "vision_sam: image.cols {} != H*W {}*{} ({})",
                image.cols,
                h,
                win,
                expected_cols
            )));
        }
        if !h.is_multiple_of(PATCH) || !win.is_multiple_of(PATCH) {
            return Err(FocrError::Other(anyhow::anyhow!(
                "vision_sam: spatial dims ({h},{win}) must be multiples of patch {PATCH}"
            )));
        }
        let gh = h / PATCH;
        let gw = win / PATCH;

        // ── patch embed: Conv2d(3->768, k16, s16), permute B,C,H,W->B,H,W,C.
        // conv2d kernel wants pre-padded NCHW; patch embed has no padding.
        let conv_out = conv_apply(&w.patch_embed, &image.data, h, win, 0, PATCH)?;
        ensure_flat_len("vision_sam patch_embed output", &conv_out, dim, gh, gw)?;
        // conv_out is [1, dim, gh, gw] (channel-major). Tokens we carry as
        // [gh*gw, dim] (spatial-major rows) for the transformer (NHWC flat).
        let mut x = nchw_to_nhwc_rows(&conv_out, dim, gh, gw);

        // ── abs pos-embed (added once before the blocks; bicubic if needed).
        let pos = abs_pos(&w.pos_embed, w.pos_grid_h, w.pos_grid_w, dim, gh, gw)?;
        for (xv, pv) in x.data.iter_mut().zip(pos.iter()) {
            *xv += *pv;
        }
        grids.push((gh, gw));
        xs.push(x);
    }

    // ── 12 transformer blocks; ONE hydration per block, all views inside. ──
    let tb = Instant::now();
    for i in 0..depth {
        let blk = block_at(i)?;
        for (x, &(gh, gw)) in xs.iter_mut().zip(grids.iter()) {
            *x = block_forward_impl(&blk, x, gh, gw, low_mem)?;
            // Sequential, on the thread that entered the forward — the only
            // place a progress event may be raised (see `super::progress`).
            super::progress::vision_step();
        }
    }
    super::timing_log(&format!(
        "    sam.blocks {:.2}s",
        tb.elapsed().as_secs_f64()
    ));

    // ── per-view neck + net_2 + net_3. ──────────────────────────────────────
    let mut out = Vec::with_capacity(xs.len());
    for (x, &(gh, gw)) in xs.iter().zip(grids.iter()) {
        // neck: x is [gh*gw, dim] NHWC rows; neck operates NCHW.
        // permute(0,3,1,2): NHWC-rows -> NCHW flat.
        let x_nchw = nhwc_rows_to_nchw(x, dim, gh, gw);

        // neck conv1: 768 -> 256, k1, no pad.
        let nc1 = conv_apply(&w.neck_conv1, &x_nchw, gh, gw, 0, 1)?;
        ensure_flat_len("vision_sam neck_conv1 output", &nc1, NECK_CH, gh, gw)?;
        let nc1 = layer_norm_2d(&nc1, &w.neck_ln1, NECK_CH, gh, gw)?;
        // neck conv2: 256 -> 256, k3, pad1.
        let nc2 = conv_apply(&w.neck_conv2, &nc1, gh, gw, 1, 1)?;
        ensure_flat_len("vision_sam neck_conv2 output", &nc2, NECK_CH, gh, gw)?;
        let neck = layer_norm_2d(&nc2, &w.neck_ln2, NECK_CH, gh, gw)?;

        // net_2: 256 -> 512, k3, s2, p1 -> grid /2.
        let (gh2, gw2) = (gh.div_ceil(2), gw.div_ceil(2));
        let x2 = conv_apply(&w.net2, &neck, gh, gw, 1, 2)?;
        ensure_flat_len("vision_sam net2 output", &x2, NET2_CH, gh2, gw2)?;
        // net_3: 512 -> 1024, k3, s2, p1 -> grid /2 again.
        let (gh3, gw3) = (gh2.div_ceil(2), gw2.div_ceil(2));
        let x3 = conv_apply(&w.net3, &x2, gh2, gw2, 1, 2)?;
        ensure_flat_len("vision_sam net3 output", &x3, OUT_CH, gh3, gw3)?;

        // x3 is [OUT_CH, gh3*gw3] NCHW flat — exactly flatten(2) layout.
        out.push(Mat::from_vec(OUT_CH, gh3 * gw3, x3));
    }
    Ok(out)
}

/// Batched SAM tower over `V` views (images) in ONE forward (bd-1azu.10).
///
/// Returns one `[OUT_CH, gh3*gw3]` matrix per view, where the `i`-th output is
/// **byte-for-byte identical** to [`forward_with`] called on `images[i]` alone.
///
/// Structure (lossless, doctrine #5 "one live forward"): every view's patch
/// tokens are stacked along a leading batch dim into a single `[V·gh·gw, dim]`
/// buffer. The transformer blocks then run their M-independent stages
/// (`norm1`/`norm2` LayerNorm and the `lin1`/`lin2` MLP linears) once over the
/// whole stack — bit-identical per row because [`nn::layer_norm`] is strictly
/// row-wise and the f32 GEMMs are M-independent. Attention is **block-diagonal**:
/// each block's [`attention`] (or window-partitioned [`attention_windowed`]) runs
/// per view over only that view's own tokens, so a view never attends across
/// into another view. The convs (patch-embed, neck, `net_2`, `net_3`, all cheap
/// relative to attention) run per-view in a sequential loop, exactly as
/// [`forward_with`] does.
///
/// No nested rayon: the outer loops over views/blocks are sequential, so each
/// `nn::matmul` / `conv_apply` / `attention` fans out one-at-a-time
/// (bd-1azu.14, doctrine #5).
///
/// # Errors
/// [`FocrError::Other`] on an empty batch, a ragged batch (a view whose
/// `image.cols != h*win`), a non-3-channel view, non-`PATCH`-divisible or
/// zero spatial dims, or any kernel rejection.
pub fn forward_with_batched(
    w: &SamWeights,
    images: &[&Mat],
    h: usize,
    win: usize,
) -> FocrResult<Vec<Mat>> {
    let v = images.len();
    if v == 0 {
        return Err(FocrError::Other(anyhow::anyhow!(
            "vision_sam forward_with_batched: empty view batch"
        )));
    }
    if h == 0 || win == 0 {
        return Err(FocrError::Other(anyhow::anyhow!(
            "vision_sam forward_with_batched: spatial dims ({h},{win}) must be non-zero"
        )));
    }
    let expected_cols = checked_shape_mul("vision_sam forward_with_batched", h, win, "H*W")?;
    if !h.is_multiple_of(PATCH) || !win.is_multiple_of(PATCH) {
        return Err(FocrError::Other(anyhow::anyhow!(
            "vision_sam forward_with_batched: spatial dims ({h},{win}) must be multiples of patch {PATCH}"
        )));
    }
    for (i, img) in images.iter().enumerate() {
        if img.rows != 3 {
            return Err(FocrError::Other(anyhow::anyhow!(
                "vision_sam forward_with_batched: view {i} expected 3 input channels, got {}",
                img.rows
            )));
        }
        if img.cols != expected_cols {
            return Err(FocrError::Other(anyhow::anyhow!(
                "vision_sam forward_with_batched: view {i} cols {} != H*W {}*{} ({expected_cols}) (ragged batch)",
                img.cols,
                h,
                win
            )));
        }
    }

    let gh = h / PATCH;
    let gw = win / PATCH;
    let n = checked_shape_mul("vision_sam forward_with_batched", gh, gw, "gh*gw")?;
    let dim = w.patch_embed.out_ch;

    // The abs pos-embed depends only on (pos_embed, gh, gw) — identical for every
    // view, so build it once and add the same contribution to each view's block.
    let pos = abs_pos(&w.pos_embed, w.pos_grid_h, w.pos_grid_w, dim, gh, gw)?;

    // ── per-view patch_embed → NHWC rows → +pos, stacked into [V*n, dim]. ──
    let row_span = checked_shape_mul("vision_sam forward_with_batched", n, dim, "n*dim")?;
    let total_rows = checked_shape_mul("vision_sam forward_with_batched", v, n, "V*n")?;
    let stacked_len = checked_shape_mul(
        "vision_sam forward_with_batched",
        total_rows,
        dim,
        "V*n*dim",
    )?;
    let mut stacked = vec![0.0f32; stacked_len];
    for (vv, img) in images.iter().enumerate() {
        let conv_out = conv_apply(&w.patch_embed, &img.data, h, win, 0, PATCH)?;
        ensure_flat_len(
            "vision_sam forward_with_batched patch_embed output",
            &conv_out,
            dim,
            gh,
            gw,
        )?;
        let x_view = nchw_to_nhwc_rows(&conv_out, dim, gh, gw);
        let base = vv * row_span;
        let dst = &mut stacked[base..base + row_span];
        for (xv, pv) in dst.iter_mut().zip(x_view.data.iter().zip(pos.iter())) {
            *xv = *pv.0 + *pv.1;
        }
    }
    let mut x = Mat::from_vec(total_rows, dim, stacked);

    // ── 12 transformer blocks over the stacked [V*n, dim] buffer. ──
    for blk in &w.blocks {
        x = block_forward_batched(blk, &x, gh, gw, v)?;
    }

    // ── split back to V per-view [n, dim], run neck + net_2 + net_3 per view. ──
    let mut out = Vec::with_capacity(v);
    for vv in 0..v {
        let base = vv * row_span;
        let x_view = Mat::from_vec(n, dim, x.data[base..base + row_span].to_vec());

        // neck: x_view is [n, dim] NHWC rows; neck operates NCHW.
        let x_nchw = nhwc_rows_to_nchw(&x_view, dim, gh, gw);
        let nc1 = conv_apply(&w.neck_conv1, &x_nchw, gh, gw, 0, 1)?;
        ensure_flat_len(
            "vision_sam forward_with_batched neck_conv1 output",
            &nc1,
            NECK_CH,
            gh,
            gw,
        )?;
        let nc1 = layer_norm_2d(&nc1, &w.neck_ln1, NECK_CH, gh, gw)?;
        let nc2 = conv_apply(&w.neck_conv2, &nc1, gh, gw, 1, 1)?;
        ensure_flat_len(
            "vision_sam forward_with_batched neck_conv2 output",
            &nc2,
            NECK_CH,
            gh,
            gw,
        )?;
        let neck = layer_norm_2d(&nc2, &w.neck_ln2, NECK_CH, gh, gw)?;

        let (gh2, gw2) = (gh.div_ceil(2), gw.div_ceil(2));
        let x2 = conv_apply(&w.net2, &neck, gh, gw, 1, 2)?;
        ensure_flat_len(
            "vision_sam forward_with_batched net2 output",
            &x2,
            NET2_CH,
            gh2,
            gw2,
        )?;
        let (gh3, gw3) = (gh2.div_ceil(2), gw2.div_ceil(2));
        let x3 = conv_apply(&w.net3, &x2, gh2, gw2, 1, 2)?;
        ensure_flat_len(
            "vision_sam forward_with_batched net3 output",
            &x3,
            OUT_CH,
            gh3,
            gw3,
        )?;
        out.push(Mat::from_vec(OUT_CH, gh3 * gw3, x3));
    }
    Ok(out)
}

// ── transformer block ──────────────────────────────────────────────────────

/// One [`BlockP`] over NHWC token rows `[gh*gw, dim]`.
///
/// `shortcut = x; x = norm1(x);` (window_partition if windowed) `x = attn(x);`
/// (window_unpartition); `x = shortcut + x; x = x + mlp(norm2(x))`.
///
/// The forwards now call [`block_forward_impl`] directly (they thread the
/// bd-4l71 `low_mem` selector); this alias pins the HISTORICAL `low_mem=false`
/// contract for the block-level unit tests.
#[cfg_attr(not(test), allow(dead_code))]
fn block_forward(blk: &BlockP, x: &Mat, gh: usize, gw: usize) -> FocrResult<Mat> {
    block_forward_impl(blk, x, gh, gw, false)
}

/// [`block_forward`] with the `low_mem` global-attention selector (bd-4l71):
/// `false` is byte-for-byte the historical path; `true` routes GLOBAL blocks
/// through the bounded-scratch [`attention_global_bounded`] (bit-identical —
/// gated by `bounded_global_attention_is_bit_identical`). Windowed blocks are
/// identical either way (their per-window logits are already tiny).
fn block_forward_impl(
    blk: &BlockP,
    x: &Mat,
    gh: usize,
    gw: usize,
    low_mem: bool,
) -> FocrResult<Mat> {
    let normed = layer_norm_rows(x, &blk.norm1)?;

    let ta = Instant::now();
    let attn_out = if blk.window > 0 {
        attention_windowed(&blk.attn, &normed, gh, gw, blk.window)?
    } else if low_mem {
        attention_global_bounded(&blk.attn, &normed, gh, gw)?
    } else {
        attention(&blk.attn, &normed, gh, gw, None)?
    };
    super::timing_log(&format!(
        "      sam.block attn({}) {:.3}s",
        if blk.window > 0 { "win" } else { "GLOBAL" },
        ta.elapsed().as_secs_f64()
    ));
    ensure_same_shape("vision_sam block attention residual", &attn_out, x)?;

    // residual 1: x = shortcut + attn
    let mut h1 = x.clone();
    for (a, b) in h1.data.iter_mut().zip(attn_out.data.iter()) {
        *a += *b;
    }

    // residual 2: x = h1 + mlp(norm2(h1))
    let tm = Instant::now();
    let normed2 = layer_norm_rows(&h1, &blk.norm2)?;
    let mlp = if low_mem {
        mlp_row_chunked(blk, &normed2)?
    } else {
        let mut hidden = blk.lin1.apply(&normed2)?;
        nn::gelu(&mut hidden);
        blk.lin2.apply(&hidden)?
    };
    super::timing_log(&format!(
        "      sam.block mlp {:.3}s",
        tm.elapsed().as_secs_f64()
    ));
    ensure_same_shape("vision_sam block mlp residual", &mlp, &h1)?;
    for (a, b) in h1.data.iter_mut().zip(mlp.data.iter()) {
        *a += *b;
    }
    Ok(h1)
}

/// Token rows the low-memory block MLP processes per pass. At the SAM tower's
/// 4096-token, 3072-wide hidden this trades ONE 48 MB `lin1` output for a
/// 6 MB scratch (the `[n, dim]` result is allocated either way); the chunk is
/// still far above the GEMM kernels' internal parallel thresholds, so the cost
/// is unmeasurable next to the residency win.
const MLP_ROW_CHUNK: usize = 512;

/// `lin2(gelu(lin1(x)))` computed in [`MLP_ROW_CHUNK`]-row passes — the
/// bd-4l71 low-memory twin of the whole-tensor MLP, so the streamed vision
/// lane never materializes the full `[n, mlp_dim]` hidden.
///
/// BIT-IDENTICAL to the whole-tensor form: every op here is row-independent —
/// [`nn::matmul`] row blocks are exact (gated by
/// `matmul_is_row_block_invariant`), the `Linear` bias add is per-row, and
/// [`nn::gelu`] is element-wise. Chunk `[r0, r1)` therefore sees exactly the
/// inputs and produces exactly the values the whole-tensor pass wrote there.
///
/// # Errors
/// Whatever [`Linear::apply`] rejects (a shape/metadata mismatch).
fn mlp_row_chunked(blk: &BlockP, normed2: &Mat) -> FocrResult<Mat> {
    mlp_row_chunked_with(blk, normed2, MLP_ROW_CHUNK)
}

/// [`mlp_row_chunked`] with an explicit chunk height, so the bit-identity gate
/// (`row_chunked_block_mlp_is_bit_identical`) can sweep partitions.
///
/// # Errors
/// Whatever [`Linear::apply`] rejects (a shape/metadata mismatch).
fn mlp_row_chunked_with(blk: &BlockP, normed2: &Mat, chunk: usize) -> FocrResult<Mat> {
    if chunk == 0 {
        return Err(FocrError::Other(anyhow::anyhow!(
            "vision_sam mlp_row_chunked: row chunk must be non-zero"
        )));
    }
    let (rows, cols) = (normed2.rows, normed2.cols);
    let mut out: Vec<f32> = Vec::with_capacity(rows * blk.lin2.out);
    let mut start = 0usize;
    while start < rows {
        let take = chunk.min(rows - start);
        let chunk = Mat::from_vec(
            take,
            cols,
            normed2.data[start * cols..(start + take) * cols].to_vec(),
        );
        let mut hidden = blk.lin1.apply(&chunk)?;
        nn::gelu(&mut hidden);
        let done = blk.lin2.apply(&hidden)?;
        out.extend_from_slice(&done.data);
        start += take;
    }
    Ok(Mat::from_vec(rows, blk.lin2.out, out))
}

/// Batched analogue of [`block_forward`] over a stacked `[V*n, dim]` buffer
/// (`n = gh*gw`), bd-1azu.10.
///
/// `norm1`, `norm2` and the MLP (`lin1`/`gelu`/`lin2`) run once over the whole
/// `V*n`-row stack — byte-identical per row to the per-view forward because
/// [`nn::layer_norm`] is row-wise, [`nn::gelu`] is element-wise, and the GEMMs in
/// [`Linear::apply`] are M-independent. Attention is computed **per view** over a
/// slice of the (row-wise-identical) normalized buffer, via the same
/// [`attention`] / [`attention_windowed`] the per-view path uses — so a view
/// never attends across into another view's tokens (block-diagonal). The two
/// residual adds are element-wise over the stack.
fn block_forward_batched(blk: &BlockP, x: &Mat, gh: usize, gw: usize, v: usize) -> FocrResult<Mat> {
    let dim = x.cols;
    let n = checked_shape_mul("vision_sam block_forward_batched", gh, gw, "gh*gw")?;
    let row_span = checked_shape_mul("vision_sam block_forward_batched", n, dim, "n*dim")?;
    let total_rows = checked_shape_mul("vision_sam block_forward_batched", v, n, "V*n")?;
    if x.rows != total_rows {
        return Err(FocrError::Other(anyhow::anyhow!(
            "vision_sam block_forward_batched: x.rows {} != V*n {total_rows}",
            x.rows
        )));
    }
    ensure_flat_len("vision_sam block_forward_batched input", &x.data, v, n, dim)?;

    // norm1 over the whole stack (row-wise -> identical to per-view).
    let normed = layer_norm_rows(x, &blk.norm1)?;

    // Attention per view (block-diagonal): slice this view's normalized rows and
    // run the exact same attention the per-view forward uses.
    let mut attn_data = vec![0.0f32; x.data.len()];
    for view in 0..v {
        let base = view * row_span;
        let normed_view = Mat::from_vec(n, dim, normed.data[base..base + row_span].to_vec());
        let av = if blk.window > 0 {
            attention_windowed(&blk.attn, &normed_view, gh, gw, blk.window)?
        } else {
            attention(&blk.attn, &normed_view, gh, gw, None)?
        };
        ensure_mat_shape(
            &av,
            n,
            dim,
            "vision_sam block_forward_batched attention view",
        )?;
        attn_data[base..base + row_span].copy_from_slice(&av.data);
    }

    // residual 1: x = shortcut + attn (element-wise over the stack).
    let mut h1 = x.clone();
    for (a, b) in h1.data.iter_mut().zip(attn_data.iter()) {
        *a += *b;
    }

    // residual 2: x = h1 + mlp(norm2(h1)), all batched over the stack.
    let normed2 = layer_norm_rows(&h1, &blk.norm2)?;
    let mut mlp = blk.lin1.apply(&normed2)?;
    nn::gelu(&mut mlp);
    let mlp = blk.lin2.apply(&mlp)?;
    ensure_same_shape("vision_sam block_forward_batched mlp residual", &mlp, &h1)?;
    for (a, b) in h1.data.iter_mut().zip(mlp.data.iter()) {
        *a += *b;
    }
    Ok(h1)
}

/// Window-partitioned multi-head attention over `norm1`-normalized NHWC token
/// rows `[gh*gw, dim]`: pad to a multiple of `ws`, run [`attention`] once per
/// `ws×ws` window (zero-padded tail tokens included exactly as upstream), then
/// `window_unpartition` back to `[gh*gw, dim]` stripping the padding.
///
/// Extracted verbatim from the windowed branch of [`block_forward`] so the
/// per-view (sequential) and batched paths share one definition — the batched
/// SAM tower (bd-1azu.10) calls this per view, guaranteeing byte-identical
/// output to the per-view forward.
fn attention_windowed(p: &AttnP, normed: &Mat, gh: usize, gw: usize, ws: usize) -> FocrResult<Mat> {
    let dim = normed.cols;
    // window_partition: pad to multiple of window, tile into win x win.
    let pad_h = (ws - gh % ws) % ws;
    let pad_w = (ws - gw % ws) % ws;
    let hp = gh + pad_h;
    let wp = gw + pad_w;
    let nwin_h = hp / ws;
    let nwin_w = wp / ws;
    let nwin = nwin_h * nwin_w;

    // Build per-window token blocks [nwin][ws*ws, dim] (zero-padded tail).
    let mut windows = vec![0.0f32; nwin * ws * ws * dim];
    for wy in 0..nwin_h {
        for wx in 0..nwin_w {
            let widx = wy * nwin_w + wx;
            for ly in 0..ws {
                for lx in 0..ws {
                    let gy = wy * ws + ly;
                    let gxx = wx * ws + lx;
                    let dst = ((widx * ws + ly) * ws + lx) * dim;
                    if gy < gh && gxx < gw {
                        let src = (gy * gw + gxx) * dim;
                        windows[dst..dst + dim].copy_from_slice(&normed.data[src..src + dim]);
                    }
                }
            }
        }
    }

    // Attention per window (each window: ws x ws grid). The rel-pos tables
    // depend only on the window size — compute them ONCE for all 25 windows
    // instead of per window (bd-av64.10).
    let nh = NUM_HEADS;
    let hd = dim / nh;
    let rh = get_rel_pos(ws, ws, &p.rel_pos_h, p.size_h, hd);
    let rw = get_rel_pos(ws, ws, &p.rel_pos_w, p.size_w, hd);
    // Windows are independent — run them across the pool (bd-av64.10: the
    // serial loop left ~90% of cores idle because each 196-token window sits
    // below the inner kernels' own parallel thresholds). Per-window
    // arithmetic is unchanged and outputs land in disjoint spans, so the
    // result is bit-identical to the serial loop.
    let win_span = ws * ws * dim;
    let mut out_windows = vec![0.0f32; windows.len()];
    out_windows
        .par_chunks_mut(win_span)
        .enumerate()
        .try_for_each(|(widx, out_chunk)| -> FocrResult<()> {
            let base = widx * win_span;
            let win_in = Mat::from_vec(ws * ws, dim, windows[base..base + win_span].to_vec());
            let win_out = attention(p, &win_in, ws, ws, Some((&rh, &rw)))?;
            out_chunk.copy_from_slice(&win_out.data);
            Ok(())
        })?;

    // window_unpartition: scatter back, strip padding.
    let mut merged = vec![0.0f32; gh * gw * dim];
    for wy in 0..nwin_h {
        for wx in 0..nwin_w {
            let widx = wy * nwin_w + wx;
            for ly in 0..ws {
                for lx in 0..ws {
                    let gy = wy * ws + ly;
                    let gxx = wx * ws + lx;
                    if gy < gh && gxx < gw {
                        let src = ((widx * ws + ly) * ws + lx) * dim;
                        let dst = (gy * gw + gxx) * dim;
                        merged[dst..dst + dim].copy_from_slice(&out_windows[src..src + dim]);
                    }
                }
            }
        }
    }
    Ok(Mat::from_vec(gh * gw, dim, merged))
}

// ── attention with decomposed relative position ([SPEC-044]) ───────────────

/// Multi-head attention over a `gh x gw` token grid `[gh*gw, dim]`, adding the
/// decomposed rel-pos bias to the logits before softmax. Returns `proj(out)`
/// shaped `[gh*gw, dim]`.
fn attention(
    p: &AttnP,
    x: &Mat,
    gh: usize,
    gw: usize,
    relpos: Option<(&[f32], &[f32])>,
) -> FocrResult<Mat> {
    let n = gh * gw;
    if x.rows != n {
        return Err(FocrError::Other(anyhow::anyhow!(
            "vision_sam attention: input rows {} != gh*gw {}*{}",
            x.rows,
            gh,
            gw
        )));
    }
    let dim = x.cols;
    let nh = NUM_HEADS;
    if dim == 0 || !dim.is_multiple_of(nh) {
        return Err(FocrError::Other(anyhow::anyhow!(
            "vision_sam attention: dim {} must be non-zero and divisible by heads {}",
            dim,
            nh
        )));
    }
    if p.qkv.in_ != dim || p.qkv.out != 3 * dim {
        return Err(FocrError::Other(anyhow::anyhow!(
            "vision_sam attention: qkv shape out/in {}x{} incompatible with dim {}",
            p.qkv.out,
            p.qkv.in_,
            dim
        )));
    }
    if p.proj.in_ != dim || p.proj.out != dim {
        return Err(FocrError::Other(anyhow::anyhow!(
            "vision_sam attention: proj shape out/in {}x{} incompatible with dim {}",
            p.proj.out,
            p.proj.in_,
            dim
        )));
    }
    let hd = dim / nh;
    ensure_rel_pos_len("rel_pos_h", p.size_h, hd, p.rel_pos_h.len())?;
    ensure_rel_pos_len("rel_pos_w", p.size_w, hd, p.rel_pos_w.len())?;
    let scale = (hd as f32).powf(-0.5);

    // qkv: [n, 3*dim].
    let qkv = p.qkv.apply(x)?;
    ensure_mat_shape(&qkv, n, 3 * dim, "vision_sam attention qkv output")?;
    // Split into per-head q,k,v: layout [n, 3, nh, hd] (PyTorch reshape order).
    // qkv row r: [3*dim] = for s in 0..3 { for head in 0..nh { hd } }.
    // We build flat per-head buffers [nh][n, hd] for q,k,v.
    let mut q = vec![0.0f32; nh * n * hd];
    let mut k = vec![0.0f32; nh * n * hd];
    let mut v = vec![0.0f32; nh * n * hd];
    for r in 0..n {
        let row = qkv.row(r);
        for head in 0..nh {
            // Each (row, head) segment is hd contiguous floats in both the
            // source row and the per-head destination — copy, don't loop.
            let dst = (head * n + r) * hd;
            q[dst..dst + hd].copy_from_slice(&row[head * hd..(head + 1) * hd]);
            k[dst..dst + hd].copy_from_slice(&row[(nh + head) * hd..(nh + head + 1) * hd]);
            v[dst..dst + hd].copy_from_slice(&row[(2 * nh + head) * hd..(2 * nh + head + 1) * hd]);
        }
    }

    // Decomposed rel-pos bias [nh][n(q), n(k)] = rel_h[qy, ky] + rel_w[qx, kx].
    // Rh = get_rel_pos(gh, gh, rel_pos_h) -> [gh, gh, hd]
    // Rw = get_rel_pos(gw, gw, rel_pos_w) -> [gw, gw, hd]
    // rel_h[head, q, ky] = sum_c q[head,q,c] * Rh[qy, ky, c]
    // rel_w[head, q, kx] = sum_c q[head,q,c] * Rw[qx, kx, c]
    //
    // The tables depend only on (gh, gw) — the windowed path passes them in,
    // computed once per block instead of once per window (bd-av64.10).
    let (rh_owned, rw_owned);
    let (rh, rw): (&[f32], &[f32]) = match relpos {
        Some((rh, rw)) => (rh, rw),
        None => {
            rh_owned = get_rel_pos(gh, gh, &p.rel_pos_h, p.size_h, hd);
            rw_owned = get_rel_pos(gw, gw, &p.rel_pos_w, p.size_w, hd);
            (&rh_owned, &rw_owned)
        }
    };

    // Compute attention head-by-head with explicit logits so we can add bias.
    // Heads are independent and write disjoint output spans — run them across
    // the pool (bd-av64.10 pass 2); per-head arithmetic is unchanged, so the
    // result is bit-identical to the sequential loop. The inner matmuls also
    // parallelize; rayon's work stealing shares the pool between the levels.
    let mut out = vec![0.0f32; nh * n * hd]; // [nh][n, hd]
    out.par_chunks_mut(n * hd)
        .enumerate()
        .try_for_each(|(head, out_chunk)| -> FocrResult<()> {
            let qh = &q[head * n * hd..(head + 1) * n * hd];
            let kh = &k[head * n * hd..(head + 1) * n * hd];
            let vh = &v[head * n * hd..(head + 1) * n * hd];
            let (rel_h_bias, rel_w_bias) = decomposed_rel_pos_bias(qh, rh, rw, gh, gw, hd);

            // logits = scale * (Q @ K^T) + decomposed rel-pos bias.
            let qh_mat = Mat::from_vec(n, hd, qh.to_vec());
            let kt_mat = Mat::from_vec(hd, n, transpose_contiguous_stores(kh, n, hd));
            let mut lm = nn::matmul(&qh_mat, &kt_mat)?;
            // Bias add, row-structured: j = ky*gw + kx, so walk (ky, kx) directly
            // instead of dividing per element — bit-identical arithmetic in the
            // same order, but branch/div-free and autovectorizable (bd-av64.10:
            // the naive j/gw + j%gw form burned ~200M integer divisions per
            // global block).
            for i in 0..n {
                let lrow = &mut lm.data[i * n..(i + 1) * n];
                let brow_w = &rel_w_bias[i * gw..(i + 1) * gw];
                for ky in 0..gh {
                    let bh = rel_h_bias[i * gh + ky];
                    let seg = &mut lrow[ky * gw..(ky + 1) * gw];
                    for (l, &bw) in seg.iter_mut().zip(brow_w) {
                        *l = scale * *l + bh + bw;
                    }
                }
            }
            // softmax rows then weighted sum of v.
            nn::softmax_rows(&mut lm)?;
            let vh_mat = Mat::from_vec(n, hd, vh.to_vec());
            let head_out = nn::matmul(&lm, &vh_mat)?;
            out_chunk.copy_from_slice(&head_out.data);
            Ok(())
        })?;

    // Reassemble [n, dim] from [nh][n, hd] (head-major -> token rows).
    let mut ctx = vec![0.0f32; n * dim];
    for head in 0..nh {
        for r in 0..n {
            let src = (head * n + r) * hd;
            let dst = r * dim + head * hd;
            ctx[dst..dst + hd].copy_from_slice(&out[src..src + hd]);
        }
    }
    let ctx_mat = Mat::from_vec(n, dim, ctx);
    let y = p.proj.apply(&ctx_mat)?;
    ensure_mat_shape(&y, n, dim, "vision_sam attention projection output")?;
    Ok(y)
}

/// Query rows per slab in [`attention_global_bounded`]: bounds the per-head
/// logits scratch to `slab × n × 4` bytes (2 MB at the 64×64 global grid)
/// instead of the full `n × n` (~64 MB × 12 concurrent heads ≈ 0.8 GB). The
/// heads run concurrently, so this slab is paid ONCE PER HEAD — 128 rows keeps
/// the whole global-attention logits footprint at ~24 MB while still handing
/// the GEMM/softmax kernels rows in the thousands of elements.
const GLOBAL_ATTN_QUERY_SLAB: usize = 128;

/// Bounded-scratch global attention (bd-4l71 wasm-lane residency): identical
/// math to [`attention`] with the per-head logits computed in CONTIGUOUS
/// query-row slabs. Each query row's logits, bias adds, softmax, and `·V`
/// reduction are computed with the SAME kernels over the SAME operands in the
/// SAME order — a query row never reduces across other query rows, so slabbing
/// only changes how rows are grouped into kernel calls. Bit-identity is gated
/// by `bounded_global_attention_is_bit_identical`.
fn attention_global_bounded(p: &AttnP, x: &Mat, gh: usize, gw: usize) -> FocrResult<Mat> {
    attention_global_bounded_with(p, x, gh, gw, GLOBAL_ATTN_QUERY_SLAB)
}

fn attention_global_bounded_with(
    p: &AttnP,
    x: &Mat,
    gh: usize,
    gw: usize,
    slab: usize,
) -> FocrResult<Mat> {
    let n = gh * gw;
    if x.rows != n {
        return Err(FocrError::Other(anyhow::anyhow!(
            "vision_sam attention: input rows {} != gh*gw {}*{}",
            x.rows,
            gh,
            gw
        )));
    }
    let dim = x.cols;
    let nh = NUM_HEADS;
    if dim == 0 || !dim.is_multiple_of(nh) {
        return Err(FocrError::Other(anyhow::anyhow!(
            "vision_sam attention: dim {} must be non-zero and divisible by heads {}",
            dim,
            nh
        )));
    }
    if p.qkv.in_ != dim || p.qkv.out != 3 * dim {
        return Err(FocrError::Other(anyhow::anyhow!(
            "vision_sam attention: qkv shape out/in {}x{} incompatible with dim {}",
            p.qkv.out,
            p.qkv.in_,
            dim
        )));
    }
    if p.proj.in_ != dim || p.proj.out != dim {
        return Err(FocrError::Other(anyhow::anyhow!(
            "vision_sam attention: proj shape out/in {}x{} incompatible with dim {}",
            p.proj.out,
            p.proj.in_,
            dim
        )));
    }
    if slab == 0 {
        return Err(FocrError::Other(anyhow::anyhow!(
            "vision_sam attention: query slab must be non-zero"
        )));
    }
    let hd = dim / nh;
    ensure_rel_pos_len("rel_pos_h", p.size_h, hd, p.rel_pos_h.len())?;
    ensure_rel_pos_len("rel_pos_w", p.size_w, hd, p.rel_pos_w.len())?;
    let scale = (hd as f32).powf(-0.5);

    // qkv + per-head split — identical to `attention`.
    let qkv = p.qkv.apply(x)?;
    ensure_mat_shape(&qkv, n, 3 * dim, "vision_sam attention qkv output")?;
    let mut q = vec![0.0f32; nh * n * hd];
    let mut k = vec![0.0f32; nh * n * hd];
    let mut v = vec![0.0f32; nh * n * hd];
    for r in 0..n {
        let row = qkv.row(r);
        for head in 0..nh {
            let dst = (head * n + r) * hd;
            q[dst..dst + hd].copy_from_slice(&row[head * hd..(head + 1) * hd]);
            k[dst..dst + hd].copy_from_slice(&row[(nh + head) * hd..(nh + head + 1) * hd]);
            v[dst..dst + hd].copy_from_slice(&row[(2 * nh + head) * hd..(2 * nh + head + 1) * hd]);
        }
    }
    // The fused `[n, 3*dim]` projection is dead once it has been split into the
    // per-head q/k/v planes; releasing it HERE (bd-4l71) keeps its ~37 MB (SAM
    // at 4096 tokens) out of the peak that the per-head parallel section below
    // sets. Pure lifetime — no value changes.
    drop(qkv);
    let rh = get_rel_pos(gh, gh, &p.rel_pos_h, p.size_h, hd);
    let rw = get_rel_pos(gw, gw, &p.rel_pos_w, p.size_w, hd);

    let mut out = vec![0.0f32; nh * n * hd]; // [nh][n, hd]
    out.par_chunks_mut(n * hd)
        .enumerate()
        .try_for_each(|(head, out_chunk)| -> FocrResult<()> {
            let qh = &q[head * n * hd..(head + 1) * n * hd];
            let kh = &k[head * n * hd..(head + 1) * n * hd];
            let vh = &v[head * n * hd..(head + 1) * n * hd];
            let (rel_h_bias, rel_w_bias) = decomposed_rel_pos_bias(qh, &rh, &rw, gh, gw, hd);
            let kt_mat = Mat::from_vec(hd, n, transpose_contiguous_stores(kh, n, hd));
            let vh_mat = Mat::from_vec(n, hd, vh.to_vec());

            // Query slabs: rows [q0, q0+qc) of the logits, biased/softmaxed/
            // reduced exactly as in `attention` (global row index i = q0 + li).
            let mut q0 = 0usize;
            while q0 < n {
                let qc = slab.min(n - q0);
                let qb_mat = Mat::from_vec(qc, hd, qh[q0 * hd..(q0 + qc) * hd].to_vec());
                let mut lm = nn::matmul(&qb_mat, &kt_mat)?;
                for li in 0..qc {
                    let i = q0 + li;
                    let lrow = &mut lm.data[li * n..(li + 1) * n];
                    let brow_w = &rel_w_bias[i * gw..(i + 1) * gw];
                    for ky in 0..gh {
                        let bh = rel_h_bias[i * gh + ky];
                        let seg = &mut lrow[ky * gw..(ky + 1) * gw];
                        for (l, &bw) in seg.iter_mut().zip(brow_w) {
                            *l = scale * *l + bh + bw;
                        }
                    }
                }
                nn::softmax_rows(&mut lm)?;
                let slab_out = nn::matmul(&lm, &vh_mat)?;
                out_chunk[q0 * hd..(q0 + qc) * hd].copy_from_slice(&slab_out.data);
                q0 += qc;
            }
            Ok(())
        })?;

    // Reassemble + project — identical to `attention`.
    let mut ctx = vec![0.0f32; n * dim];
    for head in 0..nh {
        for r in 0..n {
            let src = (head * n + r) * hd;
            let dst = r * dim + head * hd;
            ctx[dst..dst + hd].copy_from_slice(&out[src..src + hd]);
        }
    }
    let ctx_mat = Mat::from_vec(n, dim, ctx);
    let y = p.proj.apply(&ctx_mat)?;
    ensure_mat_shape(&y, n, dim, "vision_sam attention projection output")?;
    Ok(y)
}

fn ensure_mat_shape(mat: &Mat, rows: usize, cols: usize, context: &str) -> FocrResult<()> {
    if mat.rows == rows && mat.cols == cols {
        return Ok(());
    }
    Err(FocrError::Other(anyhow::anyhow!(
        "{context}: shape {:?} != expected ({rows}, {cols})",
        mat.shape()
    )))
}

fn ensure_flat_len(context: &str, data: &[f32], ch: usize, h: usize, w: usize) -> FocrResult<()> {
    let expected = checked_nchw_len(context, ch, h, w, "ch*h*w")?;
    if data.len() == expected {
        return Ok(());
    }
    Err(FocrError::Other(anyhow::anyhow!(
        "{context}: len {} != ch*h*w {}*{}*{} ({expected})",
        data.len(),
        ch,
        h,
        w
    )))
}

fn ensure_same_shape(context: &str, actual: &Mat, expected: &Mat) -> FocrResult<()> {
    if actual.shape() == expected.shape() {
        return Ok(());
    }
    Err(FocrError::Other(anyhow::anyhow!(
        "{context}: shape {:?} != expected {:?}",
        actual.shape(),
        expected.shape()
    )))
}

fn ensure_rel_pos_len(name: &str, size: usize, hd: usize, actual_len: usize) -> FocrResult<()> {
    let rows = size
        .checked_mul(2)
        .and_then(|n| n.checked_sub(1))
        .ok_or_else(|| {
            FocrError::Other(anyhow::anyhow!(
                "vision_sam attention: {name} size {size} is invalid"
            ))
        })?;
    let expected_len = rows.checked_mul(hd).ok_or_else(|| {
        FocrError::Other(anyhow::anyhow!(
            "vision_sam attention: {name} size {size} overflows for head_dim {hd}"
        ))
    })?;
    if actual_len == expected_len {
        return Ok(());
    }
    Err(FocrError::Other(anyhow::anyhow!(
        "vision_sam attention: {name} len {actual_len} != expected {expected_len} \
         for size {size} and head_dim {hd}"
    )))
}

fn decomposed_rel_pos_bias(
    qh: &[f32],
    rh: &[f32],
    rw: &[f32],
    gh: usize,
    gw: usize,
    hd: usize,
) -> (Vec<f32>, Vec<f32>) {
    let n = gh * gw;
    debug_assert_eq!(qh.len(), n * hd);
    debug_assert_eq!(rh.len(), gh * gh * hd);
    debug_assert_eq!(rw.len(), gw * gw * hd);

    let mut rel_h_bias = vec![0.0f32; n * gh];
    let mut rel_w_bias = vec![0.0f32; n * gw];
    for i in 0..n {
        let qy = i / gw;
        let qx = i % gw;
        let qi = &qh[i * hd..(i + 1) * hd];
        for ky in 0..gh {
            let rh_base = (qy * gh + ky) * hd;
            let mut bh = 0.0f32;
            for c in 0..hd {
                bh += qi[c] * rh[rh_base + c];
            }
            rel_h_bias[i * gh + ky] = bh;
        }
        for kx in 0..gw {
            let rw_base = (qx * gw + kx) * hd;
            let mut bw = 0.0f32;
            for c in 0..hd {
                bw += qi[c] * rw[rw_base + c];
            }
            rel_w_bias[i * gw + kx] = bw;
        }
    }
    (rel_h_bias, rel_w_bias)
}

/// `get_rel_pos(q_size, k_size, rel_pos)` -> `[q_size, k_size, head_dim]`.
///
/// The table `rel_pos` is `[2*size - 1, hd]`. When `size == q_size == k_size`
/// (our case — windows / global grid match the table), no interpolation is
/// needed and we index `rel_pos[(q - k) + (k_size - 1)]` directly (the
/// `q_coords - k_coords + (k_size-1)` formula with `q_size == k_size`).
fn get_rel_pos(q_size: usize, k_size: usize, rel_pos: &[f32], size: usize, hd: usize) -> Vec<f32> {
    let max_rel = 2 * q_size.max(k_size) - 1;
    // Resize the table to max_rel rows via linear interpolation if its row
    // count differs (matches F.interpolate(mode="linear") in get_rel_pos).
    let table_rows = rel_pos.len() / hd;
    debug_assert_eq!(table_rows, 2 * size - 1);
    let resized: Vec<f32> = if table_rows != max_rel {
        interp_linear_rows(rel_pos, table_rows, hd, max_rel)
    } else {
        rel_pos.to_vec()
    };

    let qf = q_size as f32;
    let kf = k_size as f32;
    let ratio_qk = (kf / qf).max(1.0);
    let ratio_kq = (qf / kf).max(1.0);
    let mut out = vec![0.0f32; q_size * k_size * hd];
    for qi in 0..q_size {
        for ki in 0..k_size {
            let qc = qi as f32 * ratio_qk;
            let kc = ki as f32 * ratio_kq;
            // PyTorch SAM uses max(q_size/k_size, 1.0) = ratio_kq for this offset
            // term (audit rank 8). Identical to ratio_qk only when q_size == k_size
            // (every current call site), so this hardens the helper for q != k.
            let rc = (qc - kc) + (k_size as f32 - 1.0) * ratio_kq;
            let idx = rc as usize; // .long() truncation
            let src = idx * hd;
            let dst = (qi * k_size + ki) * hd;
            out[dst..dst + hd].copy_from_slice(&resized[src..src + hd]);
        }
    }
    out
}

/// Linear (1-D) interpolation of a `[rows, hd]` table to `[new_rows, hd]`,
/// matching `F.interpolate(mode="linear", align_corners=False)` over the row
/// axis (per-feature). Only hit when a rel-pos table size mismatches the grid.
fn interp_linear_rows(src: &[f32], rows: usize, hd: usize, new_rows: usize) -> Vec<f32> {
    if new_rows == rows {
        return src.to_vec();
    }
    let mut out = vec![0.0f32; new_rows * hd];
    let scale = rows as f32 / new_rows as f32;
    for i in 0..new_rows {
        // align_corners=False source coordinate.
        let s = (i as f32 + 0.5) * scale - 0.5;
        let s_clamped = s.clamp(0.0, (rows - 1) as f32);
        let lo = s_clamped.floor() as usize;
        let hi = (lo + 1).min(rows - 1);
        let frac = s_clamped - lo as f32;
        for c in 0..hd {
            let a = src[lo * hd + c];
            let b = src[hi * hd + c];
            out[i * hd + c] = a + (b - a) * frac;
        }
    }
    out
}

// ── normalization helpers ──────────────────────────────────────────────────

/// LayerNorm over the last dim of token rows `[n, cols]` with affine params.
fn layer_norm_rows(x: &Mat, ln: &LayerNormP) -> FocrResult<Mat> {
    nn::layer_norm(x, Some(&ln.w), Some(&ln.b), LN_EPS)
}

/// `LayerNorm2d` over an NCHW-flat `[C, H*W]` buffer: normalize across the
/// CHANNEL axis at each spatial location, then per-channel affine
/// (`deepencoder.py:590-602`: `mean(1)`/`var(1)` over channels).
fn layer_norm_2d(
    x: &[f32],
    ln: &LayerNormP,
    ch: usize,
    gh: usize,
    gw: usize,
) -> FocrResult<Vec<f32>> {
    if ch == 0 || gh == 0 || gw == 0 {
        return Err(FocrError::Other(anyhow::anyhow!(
            "vision_sam layer_norm_2d: channels and grid dims must be non-zero (ch={ch}, gh={gh}, gw={gw})"
        )));
    }
    ensure_flat_len("vision_sam layer_norm_2d input", x, ch, gh, gw)?;
    if ln.w.len() != ch {
        return Err(FocrError::Other(anyhow::anyhow!(
            "vision_sam layer_norm_2d: weight len {} != channels {}",
            ln.w.len(),
            ch
        )));
    }
    if ln.b.len() != ch {
        return Err(FocrError::Other(anyhow::anyhow!(
            "vision_sam layer_norm_2d: bias len {} != channels {}",
            ln.b.len(),
            ch
        )));
    }
    let hw = checked_shape_mul("vision_sam layer_norm_2d", gh, gw, "gh*gw")?;
    let out_len = checked_shape_mul("vision_sam layer_norm_2d", ch, hw, "ch*gh*gw")?;
    let mut out = vec![0.0f32; out_len];
    for s in 0..hw {
        // mean over channels at spatial location s.
        let mut mean = 0.0f32;
        for c in 0..ch {
            mean += x[c * hw + s];
        }
        mean /= ch as f32;
        let mut var = 0.0f32;
        for c in 0..ch {
            let d = x[c * hw + s] - mean;
            var += d * d;
        }
        var /= ch as f32;
        let inv = 1.0 / (var + LN_EPS).sqrt();
        for c in 0..ch {
            let norm = (x[c * hw + s] - mean) * inv;
            out[c * hw + s] = ln.w[c] * norm + ln.b[c];
        }
    }
    Ok(out)
}

// ── conv + layout helpers ──────────────────────────────────────────────────
// `pub(crate)` where noted: these are the A8 shared vision conv leaves
// (bd-3jo6.1.8) — the SigLIP patch-embed (vision_siglip.rs) drives the SAME
// im2col+GEMM conv path SAM/GOT certify, exactly like GOT reuses
// `forward_prefix`/`Linear` (B3's precedent: share by import, never relocate
// certified code).

/// Apply a [`Conv`] over an NCHW-flat `[in_ch, gh*gw]` buffer with symmetric
/// zero padding `pad` and stride `stride`, returning the NCHW-flat output.
pub(crate) fn conv_apply(
    conv: &Conv,
    input: &[f32],
    gh: usize,
    gw: usize,
    pad: usize,
    stride: usize,
) -> FocrResult<Vec<f32>> {
    if gh == 0 || gw == 0 {
        return Err(FocrError::Other(anyhow::anyhow!(
            "vision_sam conv: input grid ({gh},{gw}) must be non-zero"
        )));
    }
    if stride == 0 {
        return Err(FocrError::Other(anyhow::anyhow!(
            "vision_sam conv: stride must be non-zero"
        )));
    }
    if conv.in_ch == 0 || conv.out_ch == 0 || conv.kh == 0 || conv.kw == 0 {
        return Err(FocrError::Other(anyhow::anyhow!(
            "vision_sam conv: channels and kernel dims must be non-zero (out={}, in={}, kh={}, kw={})",
            conv.out_ch,
            conv.in_ch,
            conv.kh,
            conv.kw
        )));
    }
    let expected_input =
        checked_nchw_len("vision_sam conv input", conv.in_ch, gh, gw, "in_ch*gh*gw")?;
    if input.len() != expected_input {
        return Err(FocrError::Other(anyhow::anyhow!(
            "vision_sam conv: input len {} != in_ch*gh*gw {}",
            input.len(),
            expected_input
        )));
    }
    let expected_weight = checked_conv_weight_len(
        "vision_sam conv weight",
        conv.out_ch,
        conv.in_ch,
        conv.kh,
        conv.kw,
    )?;
    if conv.w.len() != expected_weight {
        return Err(FocrError::Other(anyhow::anyhow!(
            "vision_sam conv: weight len {} != out_ch*in_ch*kh*kw {}",
            conv.w.len(),
            expected_weight
        )));
    }
    if let Some(bias) = &conv.b
        && bias.len() != conv.out_ch
    {
        return Err(FocrError::Other(anyhow::anyhow!(
            "vision_sam conv: bias len {} != out_ch {}",
            bias.len(),
            conv.out_ch
        )));
    }

    let two_pad = checked_shape_mul("vision_sam conv", 2, pad, "2*pad")?;
    let ph = checked_shape_add("vision_sam conv", gh, two_pad, "gh+2*pad")?;
    let pw = checked_shape_add("vision_sam conv", gw, two_pad, "gw+2*pad")?;
    let oh_base = checked_shape_sub("vision_sam conv", ph, conv.kh, "padded_h-kh")?;
    let ow_base = checked_shape_sub("vision_sam conv", pw, conv.kw, "padded_w-kw")?;
    let oh = checked_shape_add("vision_sam conv", oh_base / stride, 1, "output_h+1")?;
    let ow = checked_shape_add("vision_sam conv", ow_base / stride, 1, "output_w+1")?;
    let expected_out = checked_nchw_len(
        "vision_sam conv output",
        conv.out_ch,
        oh,
        ow,
        "out_ch*oh*ow",
    )?;
    let padded = pad_nchw(input, conv.in_ch, gh, gw, pad)?;
    let out = nn::conv2d(
        &padded,
        &conv.w,
        conv.b.as_deref(),
        1,
        conv.in_ch,
        ph,
        pw,
        conv.kh,
        conv.kw,
        oh,
        ow,
        stride,
        stride,
        conv.out_ch,
    );
    if out.len() != expected_out {
        return Err(FocrError::Other(anyhow::anyhow!(
            "vision_sam conv: kernel output len {} != out_ch*oh*ow {}",
            out.len(),
            expected_out
        )));
    }
    Ok(out)
}

/// Zero-pad an NCHW-flat `[ch, gh*gw]` buffer by `pad` on every spatial side ->
/// `[ch, (gh+2p)*(gw+2p)]`.
fn pad_nchw(input: &[f32], ch: usize, gh: usize, gw: usize, pad: usize) -> FocrResult<Vec<f32>> {
    let expected_input = checked_nchw_len("vision_sam pad_nchw input", ch, gh, gw, "ch*gh*gw")?;
    if input.len() != expected_input {
        return Err(FocrError::Other(anyhow::anyhow!(
            "vision_sam pad_nchw: input len {} != ch*gh*gw {}",
            input.len(),
            expected_input
        )));
    }
    if pad == 0 {
        return Ok(input.to_vec());
    }
    let two_pad = checked_shape_mul("vision_sam pad_nchw", 2, pad, "2*pad")?;
    let ph = checked_shape_add("vision_sam pad_nchw", gh, two_pad, "gh+2*pad")?;
    let pw = checked_shape_add("vision_sam pad_nchw", gw, two_pad, "gw+2*pad")?;
    let out_len = checked_nchw_len("vision_sam pad_nchw output", ch, ph, pw, "ch*ph*pw")?;
    let mut out = vec![0.0f32; out_len];
    for c in 0..ch {
        for y in 0..gh {
            for x in 0..gw {
                let src = c * gh * gw + y * gw + x;
                let dst = c * ph * pw + (y + pad) * pw + (x + pad);
                out[dst] = input[src];
            }
        }
    }
    Ok(out)
}

/// `[1, ch, gh, gw]` channel-major conv output -> NHWC token rows
/// `[gh*gw, ch]`.
pub(crate) fn nchw_to_nhwc_rows(nchw: &[f32], ch: usize, gh: usize, gw: usize) -> Mat {
    let n = gh * gw;
    let mut data = vec![0.0f32; n * ch];
    for s in 0..n {
        let dst = &mut data[s * ch..(s + 1) * ch];
        for c in 0..ch {
            dst[c] = nchw[c * n + s];
        }
    }
    Mat::from_vec(n, ch, data)
}

/// NHWC token rows `[gh*gw, ch]` -> NCHW-flat `[ch, gh*gw]`
/// (`permute(0,3,1,2)`).
fn nhwc_rows_to_nchw(x: &Mat, ch: usize, gh: usize, gw: usize) -> Vec<f32> {
    let n = gh * gw;
    debug_assert_eq!(x.rows, n);
    debug_assert_eq!(x.cols, ch);
    let mut out = vec![0.0f32; ch * n];
    for c in 0..ch {
        let dst = &mut out[c * n..(c + 1) * n];
        for (s, slot) in dst.iter_mut().enumerate() {
            *slot = x.data[s * ch + c];
        }
    }
    out
}

/// Transpose a `[rows, cols]` row-major matrix to `[cols, rows]`.
fn transpose(m: &[f32], rows: usize, cols: usize) -> Vec<f32> {
    let mut out = vec![0.0f32; rows * cols];
    for c in 0..cols {
        let dst = &mut out[c * rows..(c + 1) * rows];
        for (r, slot) in dst.iter_mut().enumerate() {
            *slot = m[r * cols + c];
        }
    }
    out
}

fn transpose_contiguous_stores(m: &[f32], rows: usize, cols: usize) -> Vec<f32> {
    let mut out = vec![0.0f32; rows * cols];
    for c in 0..cols {
        let dst = &mut out[c * rows..(c + 1) * rows];
        for r in 0..rows {
            dst[r] = m[r * cols + c];
        }
    }
    out
}

// ── absolute position embedding (bicubic interp, [SPEC-042]) ───────────────

/// Build the abs pos-embed contribution for a `[gh, gw, dim]` grid from the
/// learned `[src_h, src_w, dim]` table, bicubic-interpolating if the runtime
/// grid differs (matches `get_abs_pos_sam`). Returns NHWC token-row order
/// `[gh*gw, dim]` flattened.
fn abs_pos(
    pos: &[f32],
    src_h: usize,
    src_w: usize,
    dim: usize,
    gh: usize,
    gw: usize,
) -> FocrResult<Vec<f32>> {
    if src_h == 0 || src_w == 0 || dim == 0 || gh == 0 || gw == 0 {
        return Err(FocrError::Other(anyhow::anyhow!(
            "vision_sam abs_pos: source ({src_h},{src_w}), target ({gh},{gw}), and dim {dim} must be non-zero"
        )));
    }
    let src_hw = checked_shape_mul("vision_sam abs_pos", src_h, src_w, "src_h*src_w")?;
    let expected_pos_len = checked_shape_mul("vision_sam abs_pos", src_hw, dim, "src_h*src_w*dim")?;
    if pos.len() != expected_pos_len {
        return Err(FocrError::Other(anyhow::anyhow!(
            "vision_sam abs_pos: pos_embed len {} != src_h*src_w*dim {}*{}*{} ({expected_pos_len})",
            pos.len(),
            src_h,
            src_w,
            dim
        )));
    }
    let tgt_hw = checked_shape_mul("vision_sam abs_pos", gh, gw, "gh*gw")?;
    let out_len = checked_shape_mul("vision_sam abs_pos", tgt_hw, dim, "gh*gw*dim")?;
    if src_h == gh && src_w == gw {
        return Ok(pos.to_vec()); // already [gh*gw, dim] NHWC rows
    }
    // Interpolate per-channel over the spatial grid (bicubic, align_corners=
    // False). pos is [src_h, src_w, dim] (channel-last); we resample each
    // channel independently into [gh, gw].
    let mut out = vec![0.0f32; out_len];
    let scale_y = src_h as f32 / gh as f32;
    let scale_x = src_w as f32 / gw as f32;
    for oy in 0..gh {
        let sy = (oy as f32 + 0.5) * scale_y - 0.5;
        for ox in 0..gw {
            let sx = (ox as f32 + 0.5) * scale_x - 0.5;
            let dst = (oy * gw + ox) * dim;
            for c in 0..dim {
                out[dst + c] = bicubic_sample(pos, src_h, src_w, dim, c, sy, sx);
            }
        }
    }
    Ok(out)
}

/// Bicubic sample of channel `c` from a `[src_h, src_w, dim]` channel-last
/// table at fractional `(sy, sx)` using the Catmull-Rom-ish cubic convolution
/// kernel (`a = -0.75`, PyTorch default), edge-clamped.
fn bicubic_sample(
    pos: &[f32],
    src_h: usize,
    src_w: usize,
    dim: usize,
    c: usize,
    sy: f32,
    sx: f32,
) -> f32 {
    let iy = sy.floor();
    let ix = sx.floor();
    let fy = sy - iy;
    let fx = sx - ix;
    let wy = cubic_weights(fy);
    let wx = cubic_weights(fx);
    let mut acc = 0.0f32;
    // indexed loop: spatial kernel offset
    #[allow(clippy::needless_range_loop)]
    for m in 0..4 {
        let yy = clamp_idx(iy as isize - 1 + m as isize, src_h);
        // indexed loop: spatial kernel offset
        #[allow(clippy::needless_range_loop)]
        for n in 0..4 {
            let xx = clamp_idx(ix as isize - 1 + n as isize, src_w);
            let val = pos[(yy * src_w + xx) * dim + c];
            acc += val * wy[m] * wx[n];
        }
    }
    acc
}

/// Cubic-convolution weights for the 4-tap stencil around fractional `t` in
/// `[0,1)` with `a = -0.75` (PyTorch `mode="bicubic"`).
fn cubic_weights(t: f32) -> [f32; 4] {
    let a = -0.75f32;
    // distances of the four samples (indices -1,0,1,2) from t.
    let d0 = 1.0 + t;
    let d1 = t;
    let d2 = 1.0 - t;
    let d3 = 2.0 - t;
    [
        cubic_k(d0, a),
        cubic_k(d1, a),
        cubic_k(d2, a),
        cubic_k(d3, a),
    ]
}

/// Keys cubic kernel `W(x)` with parameter `a`.
fn cubic_k(x: f32, a: f32) -> f32 {
    let x = x.abs();
    if x <= 1.0 {
        (a + 2.0) * x * x * x - (a + 3.0) * x * x + 1.0
    } else if x < 2.0 {
        a * x * x * x - 5.0 * a * x * x + 8.0 * a * x - 4.0 * a
    } else {
        0.0
    }
}

/// Clamp an index to `[0, n-1]`.
fn clamp_idx(i: isize, n: usize) -> usize {
    if i < 0 {
        0
    } else if i as usize >= n {
        n - 1
    } else {
        i as usize
    }
}

/// Synthetic-tower construction shared by this module's tests and by the
/// model-level tests in `got`/`onechart`, which need a SAM tower to prove their
/// streamed and cached vision arms agree. Extracted rather than copied so the
/// three callers cannot drift into testing three different towers.
#[cfg(test)]
pub(crate) mod test_support {
    use super::{DEPTH, GLOBAL_BLOCKS, NUM_HEADS, PATCH, WINDOW};
    use crate::quant::focrq::{FocrqBuilder, WriteDType};

    /// Deterministic pseudo-random values in a small symmetric range. The salt
    /// is folded in AFTER the multiply and the low bits are shifted off, so
    /// nearby salts must be kept far apart or they collide into the same tensor.
    pub(crate) fn synth_values(len: usize, salt: u64) -> Vec<f32> {
        (0..len)
            .map(|i| {
                let raw = ((i as u64)
                    .wrapping_mul(6_364_136_223_846_793_005)
                    .wrapping_add(salt)
                    >> 33) as u32;
                (raw as f32 / u32::MAX as f32 - 0.5) * 0.3
            })
            .collect()
    }

    /// Write every tensor a SAM-ViT-B tower needs under `prefix`, at a reduced
    /// `dim` (must divide `NUM_HEADS`) so the test stays fast. Geometry matches
    /// the real tower's leaf names exactly, which is what lets GOT/OneChart
    /// (`model.vision_tower_high` / `model.vision_tower`) reuse it.
    pub(crate) fn add_synth_sam_tower(b: &mut FocrqBuilder, prefix: &str, dim: usize) {
        let f32_bytes =
            |values: &[f32]| -> Vec<u8> { values.iter().flat_map(|v| v.to_le_bytes()).collect() };
        let mut add = |name: String, shape: Vec<usize>, salt: u64| {
            let len: usize = shape.iter().product();
            b.add_tensor(
                name,
                WriteDType::F32,
                shape,
                f32_bytes(&synth_values(len, salt)),
            )
            .expect("valid synthetic f32 tensor");
        };

        add(
            format!("{prefix}.patch_embed.proj.weight"),
            vec![dim, 3, PATCH, PATCH],
            1,
        );
        add(format!("{prefix}.patch_embed.proj.bias"), vec![dim], 2);
        add(format!("{prefix}.pos_embed"), vec![1, 4, 4, dim], 3);
        for (idx, name) in ["neck.0", "neck.2"].iter().enumerate() {
            let ch = 8usize;
            let in_ch = if idx == 0 { dim } else { ch };
            add(
                format!("{prefix}.{name}.weight"),
                vec![ch, in_ch, 1, 1],
                10 + idx as u64,
            );
        }
        for (idx, name) in ["neck.1", "neck.3"].iter().enumerate() {
            add(format!("{prefix}.{name}.weight"), vec![8], 20 + idx as u64);
            add(format!("{prefix}.{name}.bias"), vec![8], 30 + idx as u64);
        }
        add(format!("{prefix}.net_2.weight"), vec![8, 8, 3, 3], 40);
        add(format!("{prefix}.net_3.weight"), vec![8, 8, 3, 3], 41);

        let hd = dim / NUM_HEADS;
        for i in 0..DEPTH {
            let bb = format!("{prefix}.blocks.{i}");
            let salt = 1000 * (i as u64 + 1);
            let rel_rows = if GLOBAL_BLOCKS.contains(&i) {
                2 * 4 - 1 // global grid 4
            } else {
                2 * WINDOW - 1
            };
            add(format!("{bb}.norm1.weight"), vec![dim], salt + 1);
            add(format!("{bb}.norm1.bias"), vec![dim], salt + 2);
            add(
                format!("{bb}.attn.qkv.weight"),
                vec![3 * dim, dim],
                salt + 3,
            );
            add(format!("{bb}.attn.qkv.bias"), vec![3 * dim], salt + 4);
            add(format!("{bb}.attn.proj.weight"), vec![dim, dim], salt + 5);
            add(format!("{bb}.attn.proj.bias"), vec![dim], salt + 6);
            add(format!("{bb}.attn.rel_pos_h"), vec![rel_rows, hd], salt + 7);
            add(format!("{bb}.attn.rel_pos_w"), vec![rel_rows, hd], salt + 8);
            add(format!("{bb}.norm2.weight"), vec![dim], salt + 9);
            add(format!("{bb}.norm2.bias"), vec![dim], salt + 10);
            add(
                format!("{bb}.mlp.lin1.weight"),
                vec![4 * dim, dim],
                salt + 11,
            );
            add(format!("{bb}.mlp.lin1.bias"), vec![4 * dim], salt + 12);
            add(
                format!("{bb}.mlp.lin2.weight"),
                vec![dim, 4 * dim],
                salt + 13,
            );
            add(format!("{bb}.mlp.lin2.bias"), vec![dim], salt + 14);
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::quant::focrq::{FocrqBuilder, WriteDType};
    use half::bf16;
    use serde_json::json;
    use std::time::Instant;

    /// Identity 1x1 conv (out_ch==in_ch, k=1) with an identity-ish weight is a
    /// channel mixing; here we use a diagonal weight so output == input.
    fn identity_conv1(ch: usize) -> Conv {
        let mut w = vec![0.0f32; ch * ch];
        for c in 0..ch {
            w[c * ch + c] = 1.0;
        }
        Conv {
            w,
            b: None,
            out_ch: ch,
            in_ch: ch,
            kh: 1,
            kw: 1,
        }
    }

    fn assert_err_contains<T>(res: FocrResult<T>, needle: &str) {
        let message = match res {
            Ok(_) => String::from("<ok>"),
            Err(err) => err.to_string(),
        };
        assert!(
            message.contains(needle),
            "error {message:?} did not contain {needle:?}"
        );
    }

    fn bf16_zeros(n: usize) -> Vec<u8> {
        (0..n)
            .flat_map(|_| bf16::from_f32(0.0).to_le_bytes())
            .collect()
    }

    fn add_minimal_patch_embed(b: &mut FocrqBuilder) {
        let p = "model.sam_model";
        b.add_tensor(
            format!("{p}.patch_embed.proj.weight"),
            WriteDType::Bf16,
            vec![1, 1, 1, 1],
            bf16_zeros(1),
        )
        .unwrap();
        b.add_tensor(
            format!("{p}.patch_embed.proj.bias"),
            WriteDType::Bf16,
            vec![1],
            bf16_zeros(1),
        )
        .unwrap();
    }

    // ── weight hydration error paths ───────────────────────────────────────

    #[test]
    fn sam_weights_from_rejects_rank1_patch_embed_without_panic() {
        let p = "model.sam_model";
        let mut b = FocrqBuilder::new();
        b.add_tensor(
            format!("{p}.patch_embed.proj.weight"),
            WriteDType::Bf16,
            vec![4],
            bf16_zeros(4),
        )
        .unwrap();
        let weights = Weights::from_bytes(b.build()).unwrap();
        assert_err_contains(sam_weights_from(&weights, "model.sam_model"), "rank 1");
    }

    #[test]
    fn sam_weights_from_rejects_rank1_pos_embed_without_panic() {
        let p = "model.sam_model";
        let mut b = FocrqBuilder::new();
        add_minimal_patch_embed(&mut b);
        b.add_tensor(
            format!("{p}.pos_embed"),
            WriteDType::Bf16,
            vec![4],
            bf16_zeros(4),
        )
        .unwrap();
        let weights = Weights::from_bytes(b.build()).unwrap();
        assert_err_contains(sam_weights_from(&weights, "model.sam_model"), "rank 1");
    }

    #[test]
    fn sam_weights_from_rejects_rank1_block_qkv_without_panic() {
        let p = "model.sam_model";
        let mut b = FocrqBuilder::new();
        add_minimal_patch_embed(&mut b);
        b.add_tensor(
            format!("{p}.pos_embed"),
            WriteDType::Bf16,
            vec![1, 1, 1, 1],
            bf16_zeros(1),
        )
        .unwrap();
        b.add_tensor(
            format!("{p}.blocks.0.attn.qkv.weight"),
            WriteDType::Bf16,
            vec![4],
            bf16_zeros(4),
        )
        .unwrap();
        let weights = Weights::from_bytes(b.build()).unwrap();
        assert_err_contains(sam_weights_from(&weights, "model.sam_model"), "rank 1");
    }

    #[test]
    fn transpose_roundtrips() {
        let m = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]; // [2,3]
        let t = transpose(&m, 2, 3); // [3,2]
        assert_eq!(t, vec![1.0, 4.0, 2.0, 5.0, 3.0, 6.0]);
        assert_eq!(transpose(&t, 3, 2), m);
    }

    fn test_linear(w: Vec<f32>, b: Vec<f32>, out: usize, in_: usize) -> Linear {
        Linear::from_row_major(&w, b, out, in_).expect("test linear shape is valid")
    }

    #[test]
    fn linear_applies_weight_and_bias() -> FocrResult<()> {
        // w = [[1,2,3],[4,5,6]] (out=2,in=3), b=[10,20]
        let lin = test_linear(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0], vec![10.0, 20.0], 2, 3);
        // x = [[1,1,1]] -> y = [1+2+3+10, 4+5+6+20] = [16, 35]
        let x = Mat::from_vec(1, 3, vec![1.0, 1.0, 1.0]);
        let y = lin.apply(&x)?;
        assert_eq!(y.shape(), (1, 2));
        assert!((y.data[0] - 16.0).abs() < 1e-5);
        assert!((y.data[1] - 35.0).abs() < 1e-5);
        Ok(())
    }

    #[test]
    fn linear_rejects_malformed_shapes_without_panic() {
        assert_err_contains(
            Linear::from_row_major(&[1.0; 5], vec![], 2, 3),
            "weight len",
        );
        assert_err_contains(
            Linear::from_row_major(&[1.0; 6], vec![0.0], 2, 3),
            "bias len",
        );
        let bad_x = Mat::from_vec(1, 2, vec![1.0, 2.0]);
        let lin = test_linear(vec![1.0; 6], vec![], 2, 3);
        assert_err_contains(lin.apply(&bad_x), "input cols");
    }

    #[test]
    fn pad_nchw_zeros_border() -> FocrResult<()> {
        // single channel 2x2 = [[1,2],[3,4]], pad=1 -> 4x4 with zero border.
        let input = vec![1.0, 2.0, 3.0, 4.0];
        let out = pad_nchw(&input, 1, 2, 2, 1)?;
        assert_eq!(out.len(), 16);
        // center 2x2 holds the original
        assert_eq!(out[4 + 1], 1.0);
        assert_eq!(out[4 + 2], 2.0);
        assert_eq!(out[8 + 1], 3.0);
        assert_eq!(out[8 + 2], 4.0);
        // corners are zero
        assert_eq!(out[0], 0.0);
        assert_eq!(out[15], 0.0);
        Ok(())
    }

    #[test]
    fn nchw_nhwc_roundtrip() {
        // ch=2, grid 2x2 (n=4). NCHW: c0=[1,2,3,4], c1=[5,6,7,8].
        let nchw = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0];
        let rows = nchw_to_nhwc_rows(&nchw, 2, 2, 2);
        // row 0 (spatial 0) = [c0=1, c1=5]
        assert_eq!(rows.row(0), &[1.0, 5.0]);
        assert_eq!(rows.row(3), &[4.0, 8.0]);
        let back = nhwc_rows_to_nchw(&rows, 2, 2, 2);
        assert_eq!(back, nchw);
    }

    #[test]
    fn nchw_nhwc_roundtrip_nonsquare_grid() {
        // ch=3, grid 2x3 (n=6). Each output row gathers the same spatial
        // position from all channel planes.
        let nchw = vec![
            1.0, 2.0, 3.0, 4.0, 5.0, 6.0, // c0
            10.0, 11.0, 12.0, 13.0, 14.0, 15.0, // c1
            20.0, 21.0, 22.0, 23.0, 24.0, 25.0, // c2
        ];
        let rows = nchw_to_nhwc_rows(&nchw, 3, 2, 3);
        assert_eq!(rows.row(0), &[1.0, 10.0, 20.0]);
        assert_eq!(rows.row(5), &[6.0, 15.0, 25.0]);

        let back = nhwc_rows_to_nchw(&rows, 3, 2, 3);
        assert_eq!(back, nchw);
    }

    #[test]
    fn layer_norm_2d_normalizes_channels() -> FocrResult<()> {
        // ch=2, grid 1x2 (hw=2). At spatial 0 channels=[1,3] -> mean 2, var 1.
        // normalized = [-1, 1]; affine w=[1,1], b=[0,0].
        let x = vec![1.0, 9.0, 3.0, 11.0]; // c0=[1,9], c1=[3,11] over hw=2
        let ln = LayerNormP {
            w: vec![1.0, 1.0],
            b: vec![0.0, 0.0],
        };
        let out = layer_norm_2d(&x, &ln, 2, 1, 2)?;
        // spatial 0: channels [1,3], mean 2, var 1 -> [-1, 1]
        assert!((out[0] - (-1.0)).abs() < 1e-3); // c0,s0
        assert!((out[2] - 1.0).abs() < 1e-3); // c1,s0
        // spatial 1: channels [9,11], mean 10, var 1 -> [-1, 1]
        assert!((out[1] - (-1.0)).abs() < 1e-3); // c0,s1
        assert!((out[3] - 1.0).abs() < 1e-3); // c1,s1
        Ok(())
    }

    #[test]
    fn layer_norm_2d_rejects_malformed_affine_without_panic() {
        let x = vec![1.0, 2.0, 3.0, 4.0];
        assert_err_contains(
            layer_norm_2d(
                &x,
                &LayerNormP {
                    w: vec![1.0],
                    b: vec![0.0, 0.0],
                },
                2,
                1,
                2,
            ),
            "weight len",
        );
        assert_err_contains(
            layer_norm_2d(
                &x,
                &LayerNormP {
                    w: vec![1.0, 1.0],
                    b: vec![0.0],
                },
                2,
                1,
                2,
            ),
            "bias len",
        );
    }

    #[test]
    fn layer_norm_2d_rejects_malformed_input_without_panic() {
        let ln = LayerNormP {
            w: vec![1.0, 1.0],
            b: vec![0.0, 0.0],
        };
        assert_err_contains(layer_norm_2d(&[1.0, 2.0, 3.0], &ln, 2, 1, 2), "input");
        assert_err_contains(layer_norm_2d(&[], &ln, 0, 1, 2), "non-zero");
    }

    #[test]
    fn conv_apply_identity_1x1_preserves() -> FocrResult<()> {
        // 2 channels, 2x2 grid, identity 1x1 conv -> unchanged.
        let input = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0];
        let conv = identity_conv1(2);
        let out = conv_apply(&conv, &input, 2, 2, 0, 1)?;
        assert_eq!(out, input);
        Ok(())
    }

    #[test]
    fn conv_apply_rejects_malformed_geometry_without_panic() {
        let input = vec![1.0, 2.0, 3.0, 4.0];
        let conv = identity_conv1(1);
        assert_err_contains(conv_apply(&conv, &input, 2, 2, 0, 0), "stride");

        let oversized_kernel = Conv {
            w: vec![1.0; 25],
            b: None,
            out_ch: 1,
            in_ch: 1,
            kh: 5,
            kw: 5,
        };
        assert_err_contains(
            conv_apply(&oversized_kernel, &input, 2, 2, 0, 1),
            "padded_h-kh",
        );
    }

    #[test]
    fn conv_apply_rejects_buffer_mismatches_without_panic() {
        let input = vec![1.0, 2.0, 3.0, 4.0];

        let short_weight = Conv {
            w: vec![1.0],
            b: None,
            out_ch: 2,
            in_ch: 1,
            kh: 1,
            kw: 1,
        };
        assert_err_contains(conv_apply(&short_weight, &input, 2, 2, 0, 1), "weight len");

        let bad_bias = Conv {
            w: vec![1.0],
            b: Some(vec![0.0, 1.0]),
            out_ch: 1,
            in_ch: 1,
            kh: 1,
            kw: 1,
        };
        assert_err_contains(conv_apply(&bad_bias, &input, 2, 2, 0, 1), "bias len");

        assert_err_contains(
            conv_apply(&identity_conv1(1), &input[..3], 2, 2, 0, 1),
            "input len",
        );
    }

    #[test]
    fn conv_apply_rejects_padding_overflow_before_allocating() {
        let input = vec![1.0];
        assert_err_contains(
            conv_apply(&identity_conv1(1), &input, 1, 1, usize::MAX / 2 + 1, 1),
            "2*pad",
        );
    }

    #[test]
    fn cubic_weights_sum_to_one() {
        // The cubic-convolution kernel taps sum to 1 for any fractional t.
        for &t in &[0.0f32, 0.25, 0.5, 0.75, 0.99] {
            let w = cubic_weights(t);
            let s: f32 = w.iter().sum();
            assert!((s - 1.0).abs() < 1e-5, "t={t} sum={s}");
        }
    }

    #[test]
    fn abs_pos_identity_when_grid_matches() -> FocrResult<()> {
        // src grid == target grid -> passthrough.
        let pos = vec![1.0, 2.0, 3.0, 4.0]; // 2x2x1
        let out = abs_pos(&pos, 2, 2, 1, 2, 2)?;
        assert_eq!(out, pos);
        Ok(())
    }

    #[test]
    fn abs_pos_bicubic_constant_field_is_constant() -> FocrResult<()> {
        // A constant field must remain constant under bicubic resample
        // (partition-of-unity weights). src 4x4 of 7.0 -> target 6x6.
        let dim = 1;
        let pos = vec![7.0f32; 4 * 4 * dim];
        let out = abs_pos(&pos, 4, 4, dim, 6, 6)?;
        assert_eq!(out.len(), 6 * 6);
        for &v in &out {
            assert!((v - 7.0).abs() < 1e-3, "got {v}");
        }
        Ok(())
    }

    #[test]
    fn abs_pos_rejects_malformed_source_len_without_panic() {
        assert_err_contains(abs_pos(&[1.0, 2.0, 3.0], 2, 2, 1, 2, 2), "pos_embed len");
        assert_err_contains(abs_pos(&[1.0, 2.0, 3.0], 2, 2, 1, 3, 3), "pos_embed len");
    }

    #[test]
    fn abs_pos_rejects_invalid_geometry_without_panic() {
        assert_err_contains(abs_pos(&[], 0, 2, 1, 2, 2), "must be non-zero");
        assert_err_contains(abs_pos(&[], usize::MAX, 2, 1, 2, 2), "src_h*src_w");
    }

    #[test]
    fn get_rel_pos_indexes_table_directly() {
        // size==q==k==2 -> table has 2*2-1=3 rows; hd=1.
        // table rows [10, 20, 30] for offsets (q-k)+(k-1):
        //   q0,k0: 0-0+1 = 1 -> 20
        //   q0,k1: 0-1+1 = 0 -> 10
        //   q1,k0: 1-0+1 = 2 -> 30
        //   q1,k1: 1-1+1 = 1 -> 20
        let table = vec![10.0, 20.0, 30.0];
        let r = get_rel_pos(2, 2, &table, 2, 1);
        // layout [q, k, hd]
        assert_eq!(r[0], 20.0); // q0,k0
        assert_eq!(r[1], 10.0); // q0,k1
        assert_eq!(r[2], 30.0); // q1,k0
        assert_eq!(r[3], 20.0); // q1,k1
    }

    #[test]
    fn decomposed_rel_pos_bias_matches_direct_inner_loop_formula() {
        let (gh, gw, hd) = (3, 2, 5);
        let n = gh * gw;
        let qh: Vec<f32> = (0..n * hd)
            .map(|i| ((i % 11) as f32 - 5.0) * 0.013)
            .collect();
        let rh: Vec<f32> = (0..gh * gh * hd)
            .map(|i| ((i % 7) as f32 - 3.0) * 0.017)
            .collect();
        let rw: Vec<f32> = (0..gw * gw * hd)
            .map(|i| ((i % 5) as f32 - 2.0) * 0.019)
            .collect();

        let (rel_h_bias, rel_w_bias) = decomposed_rel_pos_bias(&qh, &rh, &rw, gh, gw, hd);
        for i in 0..n {
            let qy = i / gw;
            let qx = i % gw;
            let qi = &qh[i * hd..(i + 1) * hd];
            for j in 0..n {
                let ky = j / gw;
                let kx = j % gw;
                let rh_base = (qy * gh + ky) * hd;
                let rw_base = (qx * gw + kx) * hd;
                let mut expected_h = 0.0f32;
                let mut expected_w = 0.0f32;
                for c in 0..hd {
                    expected_h += qi[c] * rh[rh_base + c];
                    expected_w += qi[c] * rw[rw_base + c];
                }
                assert_eq!(rel_h_bias[i * gh + ky], expected_h);
                assert_eq!(rel_w_bias[i * gw + kx], expected_w);
            }
        }
    }

    /// Build a tiny single-block SAM with `dim=NUM_HEADS*hd`. Here we use the
    /// real `EMBED_DIM`/`NUM_HEADS` but a tiny 2x2 grid (no window padding since
    /// global block) to exercise attention + rel-pos + residual + MLP shapes.
    fn tiny_block(window: usize) -> BlockP {
        let dim = EMBED_DIM;
        let hd = HEAD_DIM;
        // zero rel-pos so attention reduces to plain SDPA (tables sized to the
        // grid we test: 2x2 -> 2*2-1 = 3 rows).
        let size = 2;
        let rel_rows = 2 * size - 1;
        BlockP {
            norm1: LayerNormP {
                w: vec![1.0; dim],
                b: vec![0.0; dim],
            },
            attn: AttnP {
                qkv: test_linear(identity_block_3(dim), vec![0.0; 3 * dim], 3 * dim, dim),
                proj: test_linear(identity_mat(dim), vec![0.0; dim], dim, dim),
                rel_pos_h: vec![0.0; rel_rows * hd],
                rel_pos_w: vec![0.0; rel_rows * hd],
                size_h: size,
                size_w: size,
            },
            norm2: LayerNormP {
                w: vec![1.0; dim],
                b: vec![0.0; dim],
            },
            lin1: test_linear(
                vec![0.0; MLP_HIDDEN * dim],
                vec![0.0; MLP_HIDDEN],
                MLP_HIDDEN,
                dim,
            ),
            lin2: test_linear(vec![0.0; dim * MLP_HIDDEN], vec![0.0; dim], dim, MLP_HIDDEN),
            window,
        }
    }

    /// Identity `[dim, dim]` row-major.
    fn identity_mat(dim: usize) -> Vec<f32> {
        let mut w = vec![0.0f32; dim * dim];
        for i in 0..dim {
            w[i * dim + i] = 1.0;
        }
        w
    }

    /// bd-4l71 acceptance: the row-chunked block MLP (the streamed lane's
    /// bounded-scratch `lin2(gelu(lin1(x)))`) is BIT-IDENTICAL to the
    /// whole-tensor form for every row partition. Every op in the MLP is
    /// row-independent, so the chunk height can only change how many rows a
    /// kernel call sees — never a value.
    #[test]
    fn row_chunked_block_mlp_is_bit_identical() -> FocrResult<()> {
        let dim = EMBED_DIM;
        let mut blk = tiny_block(0);
        // A NON-zero MLP (tiny_block zeroes it) so the comparison has teeth.
        blk.lin1 = test_linear(
            (0..MLP_HIDDEN * dim)
                .map(|i| (((i * 31 + 7) % 101) as f32 - 50.0) * 3.0e-4)
                .collect(),
            (0..MLP_HIDDEN)
                .map(|i| ((i % 9) as f32 - 4.0) * 0.01)
                .collect(),
            MLP_HIDDEN,
            dim,
        );
        blk.lin2 = test_linear(
            (0..dim * MLP_HIDDEN)
                .map(|i| (((i * 17 + 3) % 83) as f32 - 41.0) * 4.0e-4)
                .collect(),
            (0..dim).map(|i| ((i % 6) as f32 - 3.0) * 0.02).collect(),
            dim,
            MLP_HIDDEN,
        );
        let rows = 13usize;
        let x = Mat::from_vec(
            rows,
            dim,
            (0..rows * dim)
                .map(|i| ((i % 23) as f32 - 11.0) * 0.004)
                .collect(),
        );
        let mut whole = blk.lin1.apply(&x)?;
        nn::gelu(&mut whole);
        let whole = blk.lin2.apply(&whole)?;
        for chunk in [1usize, 2, 5, 13, 512] {
            let chunked = mlp_row_chunked_with(&blk, &x, chunk)?;
            assert_eq!(chunked.shape(), whole.shape());
            assert_eq!(
                chunked
                    .data
                    .iter()
                    .map(|f| f.to_bits())
                    .collect::<Vec<u32>>(),
                whole.data.iter().map(|f| f.to_bits()).collect::<Vec<u32>>(),
                "chunk {chunk}: row-chunked MLP must be bit-identical"
            );
        }
        assert!(
            mlp_row_chunked_with(&blk, &x, 0).is_err(),
            "a zero chunk height must be rejected, not loop forever"
        );
        Ok(())
    }

    /// bd-4l71 acceptance: the bounded-scratch global attention
    /// ([`attention_global_bounded_with`]) is BIT-IDENTICAL to the historical
    /// full-logits [`attention`] across slab sizes that partition the query
    /// rows evenly, unevenly, and not at all — with a NON-identity qkv, biases,
    /// and non-zero rel-pos tables so every term is exercised.
    #[test]
    fn bounded_global_attention_is_bit_identical() -> FocrResult<()> {
        let dim = EMBED_DIM;
        let hd = HEAD_DIM;
        for &(grid, slabs) in &[(6usize, [5usize, 16, 36, 512]), (8, [3, 16, 64, 512])] {
            let n = grid * grid;
            let rel_rows = 2 * grid - 1;
            let qkv_w: Vec<f32> = (0..3 * dim * dim)
                .map(|i| (((i * 37 + 11) % 97) as f32 - 48.0) * 4.0e-4)
                .collect();
            let qkv_b: Vec<f32> = (0..3 * dim)
                .map(|i| ((i % 7) as f32 - 3.0) * 0.01)
                .collect();
            let proj_w: Vec<f32> = (0..dim * dim)
                .map(|i| (((i * 53 + 5) % 89) as f32 - 44.0) * 5.0e-4)
                .collect();
            let proj_b: Vec<f32> = (0..dim).map(|i| ((i % 5) as f32 - 2.0) * 0.02).collect();
            let attn = AttnP {
                qkv: test_linear(qkv_w, qkv_b, 3 * dim, dim),
                proj: test_linear(proj_w, proj_b, dim, dim),
                rel_pos_h: (0..rel_rows * hd)
                    .map(|i| ((i % 13) as f32 - 6.0) * 0.0011)
                    .collect(),
                rel_pos_w: (0..rel_rows * hd)
                    .map(|i| ((i % 11) as f32 - 5.0) * 0.0009)
                    .collect(),
                size_h: grid,
                size_w: grid,
            };
            let x = Mat::from_vec(
                n,
                dim,
                (0..n * dim)
                    .map(|i| ((i % 29) as f32 - 14.0) * 0.003)
                    .collect(),
            );
            let full = attention(&attn, &x, grid, grid, None)?;
            for &slab in &slabs {
                let bounded = attention_global_bounded_with(&attn, &x, grid, grid, slab)?;
                assert_eq!(bounded.shape(), full.shape());
                assert_eq!(
                    bounded
                        .data
                        .iter()
                        .map(|f| f.to_bits())
                        .collect::<Vec<u32>>(),
                    full.data.iter().map(|f| f.to_bits()).collect::<Vec<u32>>(),
                    "grid {grid} slab {slab}: bounded global attention must be bit-identical"
                );
            }
            assert!(full.data.iter().any(|&v| v != 0.0));
        }
        Ok(())
    }

    /// bd-4l71: the streamed block hydration ([`sam_block_from`]) IS the cached
    /// hydration ([`sam_weights_from`] calls it), and the streamed forward core
    /// is the cached forward core — pinned structurally here by comparing a
    /// hydrated block against the whole-tower build on synthetic weights.
    #[test]
    fn sam_block_from_matches_whole_tower_hydration() -> FocrResult<()> {
        use crate::quant::focrq::FocrqBuilder;

        let prefix = "model.sam_model";
        let dim = 24usize; // divisible by NUM_HEADS = 12
        let mut b = FocrqBuilder::new();
        // Shared with the streamed-vs-cached gate above, so both prove their
        // claim about the SAME tower.
        test_support::add_synth_sam_tower(&mut b, prefix, dim);
        let weights = Weights::from_bytes(b.build()).expect("synthetic SAM parses");

        let whole = sam_weights_from(&weights, prefix)?;
        assert_eq!(whole.blocks.len(), DEPTH);
        for i in 0..DEPTH {
            let solo = sam_block_from(&weights, prefix, i)?;
            let cached = &whole.blocks[i];
            assert_eq!(solo.window, cached.window, "block {i} window");
            assert_eq!(solo.norm1.w, cached.norm1.w, "block {i} norm1.w");
            assert_eq!(
                solo.attn.qkv.wt.data, cached.attn.qkv.wt.data,
                "block {i} qkv"
            );
            assert_eq!(solo.attn.qkv.b, cached.attn.qkv.b, "block {i} qkv bias");
            assert_eq!(
                solo.attn.proj.wt.data, cached.attn.proj.wt.data,
                "block {i} proj"
            );
            assert_eq!(
                solo.attn.rel_pos_h, cached.attn.rel_pos_h,
                "block {i} rel_pos_h"
            );
            assert_eq!(solo.lin1.wt.data, cached.lin1.wt.data, "block {i} lin1");
            assert_eq!(solo.lin2.wt.data, cached.lin2.wt.data, "block {i} lin2");
        }
        // Head-only hydration matches the whole tower's head fields.
        let head = sam_head_from(&weights, prefix)?;
        assert!(head.blocks.is_empty());
        assert_eq!(head.patch_embed.w, whole.patch_embed.w);
        assert_eq!(head.pos_embed, whole.pos_embed);
        assert_eq!(head.net3.w, whole.net3.w);
        Ok(())
    }

    /// qkv weight `[3*dim, dim]` that copies x into each of q,k,v (identity per
    /// block).
    fn identity_block_3(dim: usize) -> Vec<f32> {
        let mut w = vec![0.0f32; 3 * dim * dim];
        for s in 0..3 {
            for i in 0..dim {
                let row = s * dim + i;
                w[row * dim + i] = 1.0;
            }
        }
        w
    }

    fn attention_scalar_reference(p: &AttnP, x: &Mat, gh: usize, gw: usize) -> FocrResult<Mat> {
        let n = gh * gw;
        let dim = x.cols;
        let nh = NUM_HEADS;
        let hd = dim / nh;
        let scale = (hd as f32).powf(-0.5);
        let qkv = p.qkv.apply(x)?;
        let mut q = vec![0.0f32; nh * n * hd];
        let mut k = vec![0.0f32; nh * n * hd];
        let mut v = vec![0.0f32; nh * n * hd];
        for r in 0..n {
            let row = qkv.row(r);
            for head in 0..nh {
                for d in 0..hd {
                    q[(head * n + r) * hd + d] = row[head * hd + d];
                    k[(head * n + r) * hd + d] = row[(nh + head) * hd + d];
                    v[(head * n + r) * hd + d] = row[(2 * nh + head) * hd + d];
                }
            }
        }

        let rh = get_rel_pos(gh, gh, &p.rel_pos_h, p.size_h, hd);
        let rw = get_rel_pos(gw, gw, &p.rel_pos_w, p.size_w, hd);
        let mut out = vec![0.0f32; nh * n * hd];
        for head in 0..nh {
            let qh = &q[head * n * hd..(head + 1) * n * hd];
            let kh = &k[head * n * hd..(head + 1) * n * hd];
            let vh = &v[head * n * hd..(head + 1) * n * hd];
            let (rel_h_bias, rel_w_bias) = decomposed_rel_pos_bias(qh, &rh, &rw, gh, gw, hd);
            let mut logits = vec![0.0f32; n * n];
            for i in 0..n {
                let qi = &qh[i * hd..(i + 1) * hd];
                for j in 0..n {
                    let ky = j / gw;
                    let kx = j % gw;
                    let kj = &kh[j * hd..(j + 1) * hd];
                    let mut dot = 0.0f32;
                    for c in 0..hd {
                        dot += qi[c] * kj[c];
                    }
                    logits[i * n + j] =
                        scale * dot + rel_h_bias[i * gh + ky] + rel_w_bias[i * gw + kx];
                }
            }
            let mut lm = Mat::from_vec(n, n, logits);
            nn::softmax_rows(&mut lm)?;
            for i in 0..n {
                let probs = lm.row(i);
                let o = &mut out[(head * n + i) * hd..(head * n + i + 1) * hd];
                for (j, &pj) in probs.iter().enumerate() {
                    let vj = &vh[j * hd..(j + 1) * hd];
                    for c in 0..hd {
                        o[c] += pj * vj[c];
                    }
                }
            }
        }

        let mut ctx = vec![0.0f32; n * dim];
        for head in 0..nh {
            for r in 0..n {
                let src = (head * n + r) * hd;
                let dst = r * dim + head * hd;
                ctx[dst..dst + hd].copy_from_slice(&out[src..src + hd]);
            }
        }
        p.proj.apply(&Mat::from_vec(n, dim, ctx))
    }

    #[test]
    fn attention_gemm_matches_scalar_reference_with_relpos() -> FocrResult<()> {
        let dim = EMBED_DIM;
        let hd = HEAD_DIM;
        let grid = 3usize;
        let n = grid * grid;
        let rel_rows = 2 * grid - 1;
        let attn = AttnP {
            qkv: test_linear(identity_block_3(dim), vec![0.0; 3 * dim], 3 * dim, dim),
            proj: test_linear(identity_mat(dim), vec![0.0; dim], dim, dim),
            rel_pos_h: (0..rel_rows * hd)
                .map(|i| ((i % 13) as f32 - 6.0) * 0.0011)
                .collect(),
            rel_pos_w: (0..rel_rows * hd)
                .map(|i| ((i % 11) as f32 - 5.0) * 0.0009)
                .collect(),
            size_h: grid,
            size_w: grid,
        };
        let x = Mat::from_vec(
            n,
            dim,
            (0..n * dim)
                .map(|i| ((i % 29) as f32 - 14.0) * 0.003)
                .collect(),
        );
        let got = attention(&attn, &x, grid, grid, None)?;
        let expected = attention_scalar_reference(&attn, &x, grid, grid)?;
        let max_abs = got
            .data
            .iter()
            .zip(expected.data.iter())
            .map(|(a, b)| (a - b).abs())
            .fold(0.0f32, f32::max);
        assert!(max_abs <= 2.0e-6, "max_abs={max_abs}");
        Ok(())
    }

    #[test]
    fn attention_rejects_malformed_qkv_and_projection_shapes() {
        let dim = EMBED_DIM;
        let x = Mat::from_vec(4, dim, vec![0.0; 4 * dim]);

        // Each Linear is VALID in isolation (the constructor enforces that);
        // its out is simply incompatible with the attention dim, which is
        // exactly what attention's own shape guards must reject.
        let mut bad_qkv = tiny_block(0).attn;
        bad_qkv.qkv = test_linear(
            vec![0.0; (3 * dim - 1) * dim],
            vec![0.0; 3 * dim - 1],
            3 * dim - 1,
            dim,
        );
        assert_err_contains(attention(&bad_qkv, &x, 2, 2, None), "qkv shape");

        let mut bad_proj = tiny_block(0).attn;
        bad_proj.proj = test_linear(vec![0.0; (dim - 1) * dim], vec![0.0; dim - 1], dim - 1, dim);
        assert_err_contains(attention(&bad_proj, &x, 2, 2, None), "proj shape");
    }

    #[test]
    fn block_forward_preserves_shape_global() -> FocrResult<()> {
        // global block (window=0) on a 2x2 grid; zero MLP + zero rel-pos so
        // output = x + attn(LN(x)) with proj/qkv identity. We only assert shape
        // and that the residual path ran (output differs from input generally).
        let blk = tiny_block(0);
        let n = 4;
        let dim = EMBED_DIM;
        let mut data = vec![0.0f32; n * dim];
        for (i, v) in data.iter_mut().enumerate() {
            *v = ((i % 7) as f32) * 0.1 - 0.3;
        }
        let x = Mat::from_vec(n, dim, data);
        let out = block_forward(&blk, &x, 2, 2)?;
        assert_eq!(out.shape(), (n, dim));
        Ok(())
    }

    #[test]
    fn block_forward_windowed_pads_and_unpartitions() -> FocrResult<()> {
        // windowed block, window=3 over a 2x2 grid forces padding to 3x3 then
        // strips it back to 2x2. Shape must round-trip.
        let blk = tiny_block(3);
        let n = 4;
        let dim = EMBED_DIM;
        let data: Vec<f32> = (0..n * dim).map(|i| (i as f32 % 5.0) * 0.01).collect();
        let x = Mat::from_vec(n, dim, data);
        let out = block_forward(&blk, &x, 2, 2)?;
        assert_eq!(out.shape(), (n, dim));
        Ok(())
    }

    #[test]
    fn block_forward_rejects_mlp_residual_shape_mismatch() {
        let mut blk = tiny_block(0);
        // A VALID Linear (the constructor validates lengths) whose out is one
        // short of EMBED_DIM, so the mlp output genuinely mismatches h1 and the
        // residual shape guard fires. Mutating only the pub `out`/`b` fields
        // would leave the private pre-transposed weight at the old width and
        // panic in the bias add instead of erroring.
        blk.lin2 = test_linear(
            vec![0.0; (EMBED_DIM - 1) * MLP_HIDDEN],
            vec![0.0; EMBED_DIM - 1],
            EMBED_DIM - 1,
            MLP_HIDDEN,
        );
        let n = 4;
        let dim = EMBED_DIM;
        let data: Vec<f32> = (0..n * dim).map(|i| (i as f32 % 5.0) * 0.01).collect();
        let x = Mat::from_vec(n, dim, data);

        assert_err_contains(
            block_forward(&blk, &x, 2, 2),
            "vision_sam block mlp residual",
        );
    }

    #[test]
    fn attention_zero_relpos_is_uniform_average_for_equal_q() -> FocrResult<()> {
        // With identical token vectors, equal logits => uniform softmax =>
        // attention output == the (shared) value vector (proj identity).
        let dim = EMBED_DIM;
        let hd = HEAD_DIM;
        let size = 2;
        let rel_rows = 2 * size - 1;
        let attn = AttnP {
            qkv: test_linear(identity_block_3(dim), vec![0.0; 3 * dim], 3 * dim, dim),
            proj: test_linear(identity_mat(dim), vec![0.0; dim], dim, dim),
            rel_pos_h: vec![0.0; rel_rows * hd],
            rel_pos_w: vec![0.0; rel_rows * hd],
            size_h: size,
            size_w: size,
        };
        // all 4 tokens identical = ones
        let x = Mat::from_vec(4, dim, vec![1.0; 4 * dim]);
        let out = attention(&attn, &x, 2, 2, None)?;
        assert_eq!(out.shape(), (4, dim));
        // uniform average of identical value vectors -> 1.0 everywhere
        for &v in &out.data {
            assert!((v - 1.0).abs() < 1e-4, "got {v}");
        }
        Ok(())
    }

    #[test]
    #[ignore = "local perf probe; run explicitly with --ignored --nocapture"]
    fn sam_attention_relpos_bias_local_probe() -> FocrResult<()> {
        let dim = EMBED_DIM;
        let hd = HEAD_DIM;
        let grid = 14usize;
        let n = grid * grid;
        let rel_rows = 2 * grid - 1;
        let attn = AttnP {
            qkv: test_linear(identity_block_3(dim), vec![0.0; 3 * dim], 3 * dim, dim),
            proj: test_linear(identity_mat(dim), vec![0.0; dim], dim, dim),
            rel_pos_h: (0..rel_rows * hd)
                .map(|i| ((i % 17) as f32 - 8.0) * 0.0007)
                .collect(),
            rel_pos_w: (0..rel_rows * hd)
                .map(|i| ((i % 19) as f32 - 9.0) * 0.0005)
                .collect(),
            size_h: grid,
            size_w: grid,
        };
        let x = Mat::from_vec(
            n,
            dim,
            (0..n * dim)
                .map(|i| ((i % 31) as f32 - 15.0) * 0.002)
                .collect(),
        );

        let runs = std::env::var("FOCR_SAM_ATTN_PROBE_RUNS")
            .ok()
            .and_then(|s| s.parse::<usize>().ok())
            .unwrap_or(3)
            .max(1);
        let warm = attention(&attn, &x, grid, grid, None)?;
        let warm_checksum: f32 = warm.data.iter().step_by(97).copied().sum();
        let start = Instant::now();
        let mut checksum = 0.0f32;
        for _ in 0..runs {
            let out = attention(&attn, &x, grid, grid, None)?;
            checksum += out.data.iter().step_by(97).copied().sum::<f32>();
        }
        let elapsed = start.elapsed();
        let total_ms = elapsed.as_secs_f64() * 1000.0;
        let avg_ms = total_ms / runs as f64;
        assert!(checksum.is_finite());
        println!(
            "{}",
            json!({
                "probe": "sam_attention_relpos_bias_local_probe",
                "grid": grid,
                "tokens": n,
                "dim": dim,
                "heads": NUM_HEADS,
                "head_dim": hd,
                "runs": runs,
                "total_ms": total_ms,
                "avg_ms": avg_ms,
                "warm_checksum": warm_checksum,
                "checksum": checksum
            })
        );
        Ok(())
    }

    #[test]
    fn forward_with_end_to_end_shapes() -> FocrResult<()> {
        // Tiny end-to-end: H=W=32 image -> patch grid 2x2 -> neck keeps 2x2 ->
        // net_2 /2 -> 1x1 -> net_3 /2 -> 1x1. Output [OUT_CH, 1].
        let h = 32;
        let gh = h / PATCH; // 2
        let gw = gh;

        // patch embed conv 3->768 k16 s16
        let patch_embed = Conv {
            w: vec![0.0; EMBED_DIM * 3 * PATCH * PATCH],
            b: Some(vec![0.0; EMBED_DIM]),
            out_ch: EMBED_DIM,
            in_ch: 3,
            kh: PATCH,
            kw: PATCH,
        };

        let blocks: Vec<BlockP> = (0..DEPTH)
            .map(|i| {
                let window = if GLOBAL_BLOCKS.contains(&i) {
                    0
                } else {
                    WINDOW
                };
                // rel-pos tables must be sized to either the window or the grid.
                let size = if window == 0 { gh } else { window };
                let rel_rows = 2 * size - 1;
                let dim = EMBED_DIM;
                let hd = HEAD_DIM;
                BlockP {
                    norm1: LayerNormP {
                        w: vec![1.0; dim],
                        b: vec![0.0; dim],
                    },
                    attn: AttnP {
                        qkv: test_linear(
                            vec![0.0; 3 * dim * dim],
                            vec![0.0; 3 * dim],
                            3 * dim,
                            dim,
                        ),
                        proj: test_linear(vec![0.0; dim * dim], vec![0.0; dim], dim, dim),
                        rel_pos_h: vec![0.0; rel_rows * hd],
                        rel_pos_w: vec![0.0; rel_rows * hd],
                        size_h: size,
                        size_w: size,
                    },
                    norm2: LayerNormP {
                        w: vec![1.0; dim],
                        b: vec![0.0; dim],
                    },
                    lin1: test_linear(
                        vec![0.0; MLP_HIDDEN * dim],
                        vec![0.0; MLP_HIDDEN],
                        MLP_HIDDEN,
                        dim,
                    ),
                    lin2: test_linear(vec![0.0; dim * MLP_HIDDEN], vec![0.0; dim], dim, MLP_HIDDEN),
                    window,
                }
            })
            .collect();

        let w = SamWeights {
            patch_embed,
            pos_embed: vec![0.0; gh * gw * EMBED_DIM],
            pos_grid_h: gh,
            pos_grid_w: gw,
            blocks,
            neck_conv1: Conv {
                w: vec![0.0; NECK_CH * EMBED_DIM],
                b: None,
                out_ch: NECK_CH,
                in_ch: EMBED_DIM,
                kh: 1,
                kw: 1,
            },
            neck_ln1: LayerNormP {
                w: vec![1.0; NECK_CH],
                b: vec![0.0; NECK_CH],
            },
            neck_conv2: Conv {
                w: vec![0.0; NECK_CH * NECK_CH * 9],
                b: None,
                out_ch: NECK_CH,
                in_ch: NECK_CH,
                kh: 3,
                kw: 3,
            },
            neck_ln2: LayerNormP {
                w: vec![1.0; NECK_CH],
                b: vec![0.0; NECK_CH],
            },
            net2: Conv {
                w: vec![0.0; NET2_CH * NECK_CH * 9],
                b: None,
                out_ch: NET2_CH,
                in_ch: NECK_CH,
                kh: 3,
                kw: 3,
            },
            net3: Conv {
                w: vec![0.0; OUT_CH * NET2_CH * 9],
                b: None,
                out_ch: OUT_CH,
                in_ch: NET2_CH,
                kh: 3,
                kw: 3,
            },
        };

        let image = Mat::from_vec(3, h * h, vec![0.5; 3 * h * h]);
        let out = forward_with(&w, &image, h, h)?;
        // gh=2 -> net_2 -> 1 -> net_3 -> 1 ; 1024 channels x 1 spatial.
        assert_eq!(out.shape(), (OUT_CH, 1));
        // all-zero weights -> all-zero feature.
        assert!(out.data.iter().all(|&v| v.abs() < 1e-6));
        Ok(())
    }

    #[test]
    fn forward_with_rejects_bad_channels() {
        let w = tiny_weights_minimal();
        let bad = Mat::from_vec(2, 32 * 32, vec![0.0; 2 * 32 * 32]);
        assert!(forward_with(&w, &bad, 32, 32).is_err());
    }

    #[test]
    fn forward_with_rejects_non_patch_multiple() {
        let w = tiny_weights_minimal();
        // 20 is not a multiple of PATCH(16)
        let img = Mat::from_vec(3, 20 * 20, vec![0.0; 3 * 20 * 20]);
        assert!(forward_with(&w, &img, 20, 20).is_err());
    }

    #[test]
    fn forward_with_rejects_zero_spatial_dims_before_conv() {
        let w = tiny_weights_minimal();
        let img = Mat::from_vec(3, 0, Vec::new());
        assert!(matches!(
            forward_with(&w, &img, 0, 0),
            Err(err) if err.to_string().contains("non-zero")
        ));
    }

    #[test]
    fn forward_with_rejects_spatial_product_overflow_before_conv() {
        let w = tiny_weights_minimal();
        let img = Mat::from_vec(3, 0, Vec::new());
        assert!(matches!(
            forward_with(&w, &img, usize::MAX, 2),
            Err(err) if err.to_string().contains("H*W")
        ));
    }

    /// A structurally-valid all-zero SamWeights for negative-path tests
    /// (shape checks fire before any heavy compute).
    fn tiny_weights_minimal() -> SamWeights {
        let gh = 2;
        let blocks: Vec<BlockP> = (0..DEPTH)
            .map(|i| {
                let window = if GLOBAL_BLOCKS.contains(&i) {
                    0
                } else {
                    WINDOW
                };
                let size = if window == 0 { gh } else { window };
                let rel_rows = 2 * size - 1;
                let dim = EMBED_DIM;
                let hd = HEAD_DIM;
                BlockP {
                    norm1: LayerNormP {
                        w: vec![1.0; dim],
                        b: vec![0.0; dim],
                    },
                    attn: AttnP {
                        qkv: test_linear(
                            vec![0.0; 3 * dim * dim],
                            vec![0.0; 3 * dim],
                            3 * dim,
                            dim,
                        ),
                        proj: test_linear(vec![0.0; dim * dim], vec![0.0; dim], dim, dim),
                        rel_pos_h: vec![0.0; rel_rows * hd],
                        rel_pos_w: vec![0.0; rel_rows * hd],
                        size_h: size,
                        size_w: size,
                    },
                    norm2: LayerNormP {
                        w: vec![1.0; dim],
                        b: vec![0.0; dim],
                    },
                    lin1: test_linear(
                        vec![0.0; MLP_HIDDEN * dim],
                        vec![0.0; MLP_HIDDEN],
                        MLP_HIDDEN,
                        dim,
                    ),
                    lin2: test_linear(vec![0.0; dim * MLP_HIDDEN], vec![0.0; dim], dim, MLP_HIDDEN),
                    window,
                }
            })
            .collect();
        SamWeights {
            patch_embed: Conv {
                w: vec![0.0; EMBED_DIM * 3 * PATCH * PATCH],
                b: Some(vec![0.0; EMBED_DIM]),
                out_ch: EMBED_DIM,
                in_ch: 3,
                kh: PATCH,
                kw: PATCH,
            },
            pos_embed: vec![0.0; gh * gh * EMBED_DIM],
            pos_grid_h: gh,
            pos_grid_w: gh,
            blocks,
            neck_conv1: Conv {
                w: vec![0.0; NECK_CH * EMBED_DIM],
                b: None,
                out_ch: NECK_CH,
                in_ch: EMBED_DIM,
                kh: 1,
                kw: 1,
            },
            neck_ln1: LayerNormP {
                w: vec![1.0; NECK_CH],
                b: vec![0.0; NECK_CH],
            },
            neck_conv2: Conv {
                w: vec![0.0; NECK_CH * NECK_CH * 9],
                b: None,
                out_ch: NECK_CH,
                in_ch: NECK_CH,
                kh: 3,
                kw: 3,
            },
            neck_ln2: LayerNormP {
                w: vec![1.0; NECK_CH],
                b: vec![0.0; NECK_CH],
            },
            net2: Conv {
                w: vec![0.0; NET2_CH * NECK_CH * 9],
                b: None,
                out_ch: NET2_CH,
                in_ch: NECK_CH,
                kh: 3,
                kw: 3,
            },
            net3: Conv {
                w: vec![0.0; OUT_CH * NET2_CH * 9],
                b: None,
                out_ch: OUT_CH,
                in_ch: NET2_CH,
                kh: 3,
                kw: 3,
            },
        }
    }

    // ── bd-1azu.10: batched-SAM parity (model-free, byte-exact) ──────────────
    // Deterministic NON-trivial weights so a cross-view attention leak OR an
    // M-dependent GEMM would change the output and fail the bit-exact assert.
    // (The all-zero `tiny_weights_minimal()` is too weak — every stage collapses
    // to a constant, so it could not detect a view-to-view leak.)
    fn det(i: usize, salt: usize) -> f32 {
        let x = (i as u64)
            .wrapping_mul(2_654_435_761)
            .wrapping_add((salt as u64).wrapping_mul(40_503))
            % 1000;
        (x as f32) / 1000.0 - 0.5
    }
    fn rand_lin(out: usize, inn: usize, salt: usize) -> Linear {
        test_linear(
            (0..out * inn).map(|i| det(i, salt)).collect(),
            (0..out).map(|i| det(i, salt + 1)).collect(),
            out,
            inn,
        )
    }
    fn rand_ln(dim: usize, salt: usize) -> LayerNormP {
        LayerNormP {
            // gains near 1.0 (1 + small varied offset), varied bias.
            w: (0..dim).map(|i| 1.0 + det(i, salt)).collect(),
            b: (0..dim).map(|i| det(i, salt + 2)).collect(),
        }
    }
    fn rand_conv(
        out_ch: usize,
        in_ch: usize,
        kh: usize,
        kw: usize,
        bias: bool,
        salt: usize,
    ) -> Conv {
        Conv {
            w: (0..out_ch * in_ch * kh * kw)
                .map(|i| det(i, salt))
                .collect(),
            b: if bias {
                Some((0..out_ch).map(|i| det(i, salt + 1)).collect())
            } else {
                None
            },
            out_ch,
            in_ch,
            kh,
            kw,
        }
    }

    /// A REAL-shaped (EMBED_DIM/NUM_HEADS/DEPTH/WINDOW/GLOBAL_BLOCKS/PATCH) SAM
    /// with deterministic varied weights over a `gh×gw` patch grid. Rel-pos
    /// tables are sized to the window (windowed blocks) or the grid (global
    /// blocks), so no rel-pos interpolation fires (the deployed identity path).
    fn tiny_weights_nontrivial(gh: usize, gw: usize) -> SamWeights {
        let dim = EMBED_DIM;
        let hd = HEAD_DIM;
        let blocks: Vec<BlockP> = (0..DEPTH)
            .map(|i| {
                let window = if GLOBAL_BLOCKS.contains(&i) {
                    0
                } else {
                    WINDOW
                };
                let size_h = if window == 0 { gh } else { window };
                let size_w = if window == 0 { gw } else { window };
                let rel_rows_h = 2 * size_h - 1;
                let rel_rows_w = 2 * size_w - 1;
                BlockP {
                    norm1: rand_ln(dim, 100 + i * 13),
                    attn: AttnP {
                        qkv: rand_lin(3 * dim, dim, 200 + i * 13),
                        proj: rand_lin(dim, dim, 300 + i * 13),
                        rel_pos_h: (0..rel_rows_h * hd).map(|j| det(j, 700 + i * 13)).collect(),
                        rel_pos_w: (0..rel_rows_w * hd).map(|j| det(j, 800 + i * 13)).collect(),
                        size_h,
                        size_w,
                    },
                    norm2: rand_ln(dim, 400 + i * 13),
                    lin1: rand_lin(MLP_HIDDEN, dim, 500 + i * 13),
                    lin2: rand_lin(dim, MLP_HIDDEN, 600 + i * 13),
                    window,
                }
            })
            .collect();
        SamWeights {
            patch_embed: rand_conv(EMBED_DIM, 3, PATCH, PATCH, true, 11),
            pos_embed: (0..gh * gw * EMBED_DIM).map(|i| det(i, 13)).collect(),
            pos_grid_h: gh,
            pos_grid_w: gw,
            blocks,
            neck_conv1: rand_conv(NECK_CH, EMBED_DIM, 1, 1, false, 21),
            neck_ln1: rand_ln(NECK_CH, 23),
            neck_conv2: rand_conv(NECK_CH, NECK_CH, 3, 3, false, 25),
            neck_ln2: rand_ln(NECK_CH, 27),
            net2: rand_conv(NET2_CH, NECK_CH, 3, 3, false, 29),
            net3: rand_conv(OUT_CH, NET2_CH, 3, 3, false, 31),
        }
    }

    /// A `[3, h*win]` image with deterministic varied (bounded) pixels.
    fn nontrivial_image(h: usize, win: usize, salt: usize) -> Mat {
        let cols = h * win;
        Mat::from_vec(
            3,
            cols,
            (0..3 * cols)
                .map(|i| ((((i + salt * 7919) as f32) * 0.0007).sin()) * 0.5)
                .collect(),
        )
    }

    #[test]
    fn batched_sam_equals_per_view_byte_for_byte() -> FocrResult<()> {
        // H=W=240 -> gh=gw=15: windowed blocks (WINDOW=14) partition into a 2x2
        // grid of windows (real + zero-padded tail) AND global blocks run on the
        // 15x15 grid -> a cross-view leak in EITHER path would change a byte.
        let h = 240;
        let win = 240;
        let gh = h / PATCH;
        let gw = win / PATCH;
        let w = tiny_weights_nontrivial(gh, gw);
        let v0 = nontrivial_image(h, win, 1);
        let v1 = nontrivial_image(h, win, 2);
        let v2 = nontrivial_image(h, win, 3);

        let batched = forward_with_batched(&w, &[&v0, &v1, &v2], h, win)?;
        assert_eq!(batched.len(), 3);
        for (i, view) in [&v0, &v1, &v2].iter().enumerate() {
            let seq = forward_with(&w, view, h, win)?;
            assert_eq!(batched[i].shape(), seq.shape(), "view {i} shape");
            assert_eq!(
                batched[i].data, seq.data,
                "view {i}: batched SAM != per-view sequential (cross-view leak or M-dependence)"
            );
        }
        Ok(())
    }

    #[test]
    fn batched_sam_single_view_equals_forward_with() -> FocrResult<()> {
        let h = 240;
        let win = 240;
        let gh = h / PATCH;
        let gw = win / PATCH;
        let w = tiny_weights_nontrivial(gh, gw);
        let view = nontrivial_image(h, win, 5);
        let batched = forward_with_batched(&w, &[&view], h, win)?;
        let seq = forward_with(&w, &view, h, win)?;
        assert_eq!(batched.len(), 1);
        assert_eq!(batched[0].shape(), seq.shape());
        assert_eq!(batched[0].data, seq.data);
        Ok(())
    }

    /// bd-K2 acceptance: hoisting the streamed lane's per-block hydration OUT
    /// of the view loop (views-inner [`forward_core_views`]) reproduces the
    /// views-outer nest ([`forward_core`], one hydration sweep per view)
    /// BIT-FOR-BIT. `block_at` hands out an OWNED clone per request, exactly as
    /// [`forward_streamed_views`] hands out a freshly dequantized block, and
    /// `low_mem = true` selects the same bounded global-attention kernel the
    /// streamed lane uses.
    #[test]
    fn streamed_views_inner_equals_views_outer_byte_for_byte() -> FocrResult<()> {
        let h = 240;
        let win = 240;
        let gh = h / PATCH;
        let gw = win / PATCH;
        let w = tiny_weights_nontrivial(gh, gw);
        // The "streamed" head: block-free, blocks arrive from `block_at`.
        let mut head = tiny_weights_nontrivial(gh, gw);
        head.blocks = Vec::new();
        let blocks = w.blocks.clone();
        let mut hydrate = |i: usize| -> FocrResult<std::borrow::Cow<'static, BlockP>> {
            Ok(std::borrow::Cow::Owned(blocks[i].clone()))
        };

        // Two same-size views (the deployed unlimited page) plus a smaller one,
        // so per-view geometry (rel-pos interpolation on the global blocks) is
        // exercised too.
        let small = 128usize;
        let views = [
            nontrivial_image(h, win, 1),
            nontrivial_image(h, win, 2),
            nontrivial_image(small, small, 3),
        ];
        let dims = [(h, win), (h, win), (small, small)];

        // views-outer: a full hydration sweep per view (the old shape).
        let mut per_view = Vec::new();
        for (view, &(vh, vw)) in views.iter().zip(dims.iter()) {
            per_view.push(forward_core(
                &head,
                view,
                vh,
                vw,
                DEPTH,
                &mut hydrate,
                true,
            )?);
        }

        // views-inner: ONE hydration sweep, every view through each block.
        let refs: Vec<&Mat> = views.iter().collect();
        let hoisted = forward_core_views(&head, &refs, &dims, DEPTH, &mut hydrate, true)?;

        assert_eq!(hoisted.len(), per_view.len());
        for (i, (outer, inner)) in per_view.iter().zip(hoisted.iter()).enumerate() {
            assert_eq!(outer.shape(), inner.shape(), "view {i} shape");
            assert_eq!(
                outer.data.iter().map(|f| f.to_bits()).collect::<Vec<u32>>(),
                inner.data.iter().map(|f| f.to_bits()).collect::<Vec<u32>>(),
                "view {i}: hydration-hoisted SAM must be bit-identical to per-view"
            );
        }
        // Non-degeneracy + the views really differ (a cross-view leak or a
        // stale activation would otherwise pass trivially).
        assert!(per_view[0].data.iter().any(|&v| v != 0.0));
        assert_ne!(per_view[0].data, per_view[1].data);
        Ok(())
    }

    #[test]
    fn batched_sam_rejects_ragged_and_empty() {
        let h = 32;
        let win = 32;
        let w = tiny_weights_minimal(); // validation fires before any compute
        // empty batch -> Err.
        assert!(forward_with_batched(&w, &[], h, win).is_err());
        // ragged: second view has a different H*W (cols mismatch) -> Err.
        let a = nontrivial_image(h, win, 1);
        let b = nontrivial_image(h, win + PATCH, 2); // cols = h*(win+16) != h*win
        assert!(forward_with_batched(&w, &[&a, &b], h, win).is_err());
        // non-3-channel view -> Err.
        let bad = Mat::from_vec(2, h * win, vec![0.0; 2 * h * win]);
        assert!(forward_with_batched(&w, &[&bad], h, win).is_err());
    }
}