environment/training-source/source-snapshot/examples/verify.rs

//! Compare our encoder against the HuggingFace reference, on desktop.
//!
//! ViT numerics are unforgiving: a transposed weight, a wrong RoPE pairing,
//! or a misordered patch flattening all produce output that looks
//! statistically reasonable and is completely wrong. This has to pass
//! before anything is deployed, because diagnosing it through logcat is
//! miserable.
//!
//! Inputs come from `tools/dump_reference.py`, which writes the *already
//! preprocessed* pixel tensor alongside the expected features. Taking the
//! pixel tensor verbatim keeps image resizing out of the comparison, so a
//! failure here is a failure in the graph.
//!
//! ```text
//! python tools/dump_reference.py --out ref/
//! cargo run --release --bin verify -- ref/ [model.safetensors]
//! ```

use std::path::{Path, PathBuf};

use dinovision::dinov3::Config;
use meganeura::Graph;
use meganeura::train::{Mode, SessionConfig};
use serde::{Deserialize, Serialize};

mod common;

#[derive(Serialize)]
struct Verification {
    schema_version: u32,
    image_size: usize,
    encoder_layers: usize,
    model_sha256: String,
    source_model_sha256: Option<String>,
    pixel_values_sha256: String,
    reference_features_sha256: String,
    meganeura_features_sha256: String,
    elements: usize,
    relative_l2: f64,
    max_absolute_error: f32,
    worst_absolute_token: usize,
    cls_cosine: f32,
    worst_patch_cosine: f32,
    stages: Vec<StageMetric>,
    thresholds: Thresholds,
    passed: bool,
}

#[derive(Deserialize)]
struct ReferenceMetadata {
    image_size: usize,
    encoder_layers: usize,
    pixel_values_sha256: String,
    features_sha256: String,
    exported_model_sha256: String,
    source_model_sha256: Option<String>,
}

#[derive(Serialize)]
struct StageMetric {
    name: String,
    reference_file: String,
    meganeura_file: String,
    elements: usize,
    reference_sha256: String,
    meganeura_sha256: String,
    relative_l2: f64,
    max_absolute_error: f32,
    cosine: f32,
}

#[derive(Serialize)]
struct Thresholds {
    max_relative_l2: f64,
    min_cls_cosine: f32,
    min_patch_cosine: f32,
}

fn read_f32(path: &Path) -> Vec<f32> {
    let bytes = std::fs::read(path).unwrap_or_else(|e| panic!("{}: {e}", path.display()));
    assert_eq!(
        bytes.len() % 4,
        0,
        "{} is not a whole number of f32 values",
        path.display()
    );
    bytes
        .chunks_exact(4)
        .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]]))
        .collect()
}

fn embedding_forward(
    gpu: std::sync::Arc<blade_graphics::Context>,
    config: &Config,
    model: &meganeura::data::safetensors::SafeTensorsModel,
    patches: &[f32],
) -> Vec<f32> {
    let hidden = config.hidden_size;
    let prefix = config.num_prefix_tokens();
    let mut graph = Graph::new();
    let input = graph.input("patches", &[config.num_patches(), config.patch_dim()]);
    let weight = graph.parameter(
        "embeddings.patch_embeddings.weight",
        &[config.patch_dim(), hidden],
    );
    let bias = graph.parameter("embeddings.patch_embeddings.bias", &[hidden]);
    let patch_embeddings = graph.matmul(input, weight);
    let patch_embeddings = graph.bias_add(patch_embeddings, bias);
    let prefix_tokens = graph.parameter("prefix_tokens", &[prefix, hidden]);
    let output = graph.concat(
        prefix_tokens,
        patch_embeddings,
        1,
        (prefix * hidden) as u32,
        (config.num_patches() * hidden) as u32,
        1,
    );
    let output = graph.reshape(output, &[config.num_tokens(), hidden]);
    graph.set_outputs(vec![output]);

    let (mut session, _) = meganeura::train::build(
        &graph,
        SessionConfig {
            mode: Mode::Inference,
            gpu: Some(gpu),
            ..Default::default()
        },
    );
    let convolution = model
        .tensor_f32_auto("embeddings.patch_embeddings.weight")
        .expect("load patch embedding weight");
    session.set_parameter(
        "embeddings.patch_embeddings.weight",
        &dinovision::preprocess::conv_weight_to_matmul(&convolution, hidden, config.patch_dim()),
    );
    session.set_parameter(
        "embeddings.patch_embeddings.bias",
        &model
            .tensor_f32_auto("embeddings.patch_embeddings.bias")
            .expect("load patch embedding bias"),
    );
    let mut prefix_data = model
        .tensor_f32_auto("embeddings.cls_token")
        .expect("load CLS token");
    prefix_data.extend(
        model
            .tensor_f32_auto("embeddings.register_tokens")
            .expect("load register tokens"),
    );
    session.set_parameter("prefix_tokens", &prefix_data);
    session.set_input("patches", patches);
    session.step();
    session.wait();
    session.read_output(config.num_tokens() * hidden)
}

fn first_projections(
    gpu: std::sync::Arc<blade_graphics::Context>,
    config: &Config,
    model: &meganeura::data::safetensors::SafeTensorsModel,
    embeddings: &[f32],
) -> [Vec<f32>; 4] {
    let hidden = config.hidden_size;
    let mut graph = Graph::new();
    let input = graph.input("embeddings", &[config.num_tokens(), hidden]);
    let norm_weight = graph.parameter("layer.0.norm1.weight", &[hidden]);
    let norm_bias = graph.parameter("layer.0.norm1.bias", &[hidden]);
    let normalized = graph.layer_norm(input, norm_weight, norm_bias, config.layer_norm_eps);
    let mut projections = Vec::new();
    for (name, has_bias) in [("q_proj", true), ("k_proj", false), ("v_proj", true)] {
        let weight = graph.parameter(
            &format!("layer.0.attention.{name}.weight"),
            &[hidden, hidden],
        );
        let projected = graph.matmul(normalized, weight);
        let projected = if has_bias {
            let bias = graph.parameter(&format!("layer.0.attention.{name}.bias"), &[hidden]);
            graph.bias_add(projected, bias)
        } else {
            projected
        };
        projections.push(projected);
    }
    graph.set_outputs(vec![
        normalized,
        projections[0],
        projections[1],
        projections[2],
    ]);
    let (mut session, _) = meganeura::train::build(
        &graph,
        SessionConfig {
            mode: Mode::Inference,
            gpu: Some(gpu),
            ..Default::default()
        },
    );
    let prefix = if model
        .tensor_info()
        .contains_key("model.layer.0.norm1.weight")
    {
        "model."
    } else {
        ""
    };
    for part in ["weight", "bias"] {
        session.set_parameter(
            &format!("layer.0.norm1.{part}"),
            &model
                .tensor_f32_auto(&format!("{prefix}layer.0.norm1.{part}"))
                .expect("load first norm"),
        );
    }
    for (name, has_bias) in [("q_proj", true), ("k_proj", false), ("v_proj", true)] {
        session.set_parameter(
            &format!("layer.0.attention.{name}.weight"),
            &model
                .tensor_f32_auto_transposed(&format!("{prefix}layer.0.attention.{name}.weight"))
                .expect("load first projection"),
        );
        if has_bias {
            session.set_parameter(
                &format!("layer.0.attention.{name}.bias"),
                &model
                    .tensor_f32_auto(&format!("{prefix}layer.0.attention.{name}.bias"))
                    .expect("load first projection bias"),
            );
        }
    }
    session.set_input("embeddings", embeddings);
    session.step();
    session.wait();
    std::array::from_fn(|index| {
        let mut output = vec![0.0; config.num_tokens() * hidden];
        session.read_output_by_index(index, &mut output);
        output
    })
}

fn first_attention(
    gpu: std::sync::Arc<blade_graphics::Context>,
    config: &Config,
    q: &[f32],
    k: &[f32],
    v: &[f32],
) -> [Vec<f32>; 3] {
    let tokens = config.num_tokens();
    let heads = config.num_attention_heads;
    let head_dim = config.head_dim();
    let hidden = config.hidden_size;
    let mut graph = Graph::new();
    let q_node = graph.input("q", &[tokens, hidden]);
    let k_node = graph.input("k", &[tokens, hidden]);
    let v_node = graph.input("v", &[tokens, hidden]);
    let (cos_data, sin_data) = dinovision::dinov3::rope_tables(config);
    let cos = graph.constant(cos_data, &[tokens, hidden]);
    let sin = graph.constant(sin_data, &[tokens, hidden]);
    let apply_rope = |graph: &mut Graph, input| {
        let blocks = tokens as u32 * heads;
        let half = head_dim / 2;
        let first = graph.split_a(input, blocks, half, half, 1);
        let second = graph.split_b(input, blocks, half, half, 1);
        let negative_second = graph.neg(second);
        let rotated = graph.concat(negative_second, first, blocks, half, half, 1);
        let rotated = graph.reshape(rotated, &[tokens, hidden]);
        let straight = graph.mul(input, cos);
        let crossed = graph.mul(rotated, sin);
        graph.add(straight, crossed)
    };
    let q_rope = apply_rope(&mut graph, q_node);
    let k_rope = apply_rope(&mut graph, k_node);
    let attention = graph.full_attention(q_rope, k_rope, v_node, heads, heads, head_dim);
    graph.set_outputs(vec![q_rope, k_rope, attention]);
    let (mut session, _) = meganeura::train::build(
        &graph,
        SessionConfig {
            mode: Mode::Inference,
            gpu: Some(gpu),
            ..Default::default()
        },
    );
    session.set_input("q", q);
    session.set_input("k", k);
    session.set_input("v", v);
    session.step();
    session.wait();
    std::array::from_fn(|index| {
        let mut output = vec![0.0; tokens * hidden];
        session.read_output_by_index(index, &mut output);
        output
    })
}

fn first_remainder(
    gpu: std::sync::Arc<blade_graphics::Context>,
    config: &Config,
    model: &meganeura::data::safetensors::SafeTensorsModel,
    embeddings: &[f32],
    attention: &[f32],
) -> Vec<Vec<f32>> {
    let tokens = config.num_tokens();
    let hidden = config.hidden_size;
    let intermediate = config.intermediate_size;
    let mut graph = Graph::new();
    let embeddings_node = graph.input("embeddings", &[tokens, hidden]);
    let attention_node = graph.input("attention", &[tokens, hidden]);
    let output_weight = graph.parameter("layer.0.attention.o_proj.weight", &[hidden, hidden]);
    let output_bias = graph.parameter("layer.0.attention.o_proj.bias", &[hidden]);
    let attention_projected = graph.matmul(attention_node, output_weight);
    let attention_projected = graph.bias_add(attention_projected, output_bias);
    let scale1 = graph.parameter("layer.0.layer_scale1.lambda1", &[hidden]);
    let transposed = graph.transpose(attention_projected);
    let flat = graph.reshape(transposed, &[hidden * tokens]);
    let scaled = graph.mul_per_channel(flat, scale1, hidden as u32, tokens as u32);
    let scaled = graph.reshape(scaled, &[hidden, tokens]);
    let attention_scaled = graph.transpose(scaled);
    let residual = graph.add(embeddings_node, attention_scaled);

    let norm2_weight = graph.parameter("layer.0.norm2.weight", &[hidden]);
    let norm2_bias = graph.parameter("layer.0.norm2.bias", &[hidden]);
    let norm2 = graph.layer_norm(residual, norm2_weight, norm2_bias, config.layer_norm_eps);
    let up_weight = graph.parameter("layer.0.mlp.up_proj.weight", &[hidden, intermediate]);
    let up_bias = graph.parameter("layer.0.mlp.up_proj.bias", &[intermediate]);
    let mlp_up = graph.matmul(norm2, up_weight);
    let mlp_up = graph.bias_add(mlp_up, up_bias);
    let mlp_activated = graph.gelu(mlp_up);
    let down_weight = graph.parameter("layer.0.mlp.down_proj.weight", &[intermediate, hidden]);
    let down_bias = graph.parameter("layer.0.mlp.down_proj.bias", &[hidden]);
    let mlp_down = graph.matmul(mlp_activated, down_weight);
    let mlp_down = graph.bias_add(mlp_down, down_bias);
    let scale2 = graph.parameter("layer.0.layer_scale2.lambda1", &[hidden]);
    let transposed = graph.transpose(mlp_down);
    let flat = graph.reshape(transposed, &[hidden * tokens]);
    let scaled = graph.mul_per_channel(flat, scale2, hidden as u32, tokens as u32);
    let scaled = graph.reshape(scaled, &[hidden, tokens]);
    let mlp_scaled = graph.transpose(scaled);
    let layer_output = graph.add(residual, mlp_scaled);
    let final_weight = graph.parameter("norm.weight", &[hidden]);
    let final_bias = graph.parameter("norm.bias", &[hidden]);
    let final_norm = graph.layer_norm(
        layer_output,
        final_weight,
        final_bias,
        config.layer_norm_eps,
    );
    graph.set_outputs(vec![
        attention_projected,
        attention_scaled,
        residual,
        norm2,
        mlp_up,
        mlp_activated,
        mlp_down,
        mlp_scaled,
        layer_output,
        final_norm,
    ]);

    let (mut session, _) = meganeura::train::build(
        &graph,
        SessionConfig {
            mode: Mode::Inference,
            gpu: Some(gpu),
            ..Default::default()
        },
    );
    let prefix = if model
        .tensor_info()
        .contains_key("model.layer.0.norm1.weight")
    {
        "model."
    } else {
        ""
    };
    for name in [
        "attention.o_proj.weight",
        "mlp.up_proj.weight",
        "mlp.down_proj.weight",
    ] {
        session.set_parameter(
            &format!("layer.0.{name}"),
            &model
                .tensor_f32_auto_transposed(&format!("{prefix}layer.0.{name}"))
                .expect("load first-layer matrix"),
        );
    }
    for name in [
        "attention.o_proj.bias",
        "layer_scale1.lambda1",
        "norm2.weight",
        "norm2.bias",
        "mlp.up_proj.bias",
        "mlp.down_proj.bias",
        "layer_scale2.lambda1",
    ] {
        session.set_parameter(
            &format!("layer.0.{name}"),
            &model
                .tensor_f32_auto(&format!("{prefix}layer.0.{name}"))
                .expect("load first-layer vector"),
        );
    }
    for part in ["weight", "bias"] {
        session.set_parameter(
            &format!("norm.{part}"),
            &model
                .tensor_f32_auto(&format!("norm.{part}"))
                .expect("load final norm"),
        );
    }
    session.set_input("embeddings", embeddings);
    session.set_input("attention", attention);
    session.step();
    session.wait();
    [
        hidden,
        hidden,
        hidden,
        hidden,
        intermediate,
        intermediate,
        hidden,
        hidden,
        hidden,
        hidden,
    ]
    .into_iter()
    .enumerate()
    .map(|(index, width)| {
        let mut output = vec![0.0; tokens * width];
        session.read_output_by_index(index, &mut output);
        output
    })
    .collect()
}

fn relative_l2(actual: &[f32], expected: &[f32]) -> f64 {
    assert_eq!(actual.len(), expected.len());
    let squared_error: f64 = actual
        .iter()
        .zip(expected)
        .map(|(a, b)| (*a as f64 - *b as f64).powi(2))
        .sum();
    let squared_reference: f64 = expected.iter().map(|x| (*x as f64).powi(2)).sum();
    squared_error.sqrt() / squared_reference.sqrt().max(f64::MIN_POSITIVE)
}

fn cosine(actual: &[f32], expected: &[f32]) -> f32 {
    assert_eq!(actual.len(), expected.len());
    let dot: f64 = actual
        .iter()
        .zip(expected)
        .map(|(a, b)| *a as f64 * *b as f64)
        .sum();
    let actual_norm = actual
        .iter()
        .map(|value| (*value as f64).powi(2))
        .sum::<f64>()
        .sqrt();
    let expected_norm = expected
        .iter()
        .map(|value| (*value as f64).powi(2))
        .sum::<f64>()
        .sqrt();
    (dot / (actual_norm * expected_norm).max(f64::MIN_POSITIVE)) as f32
}

fn compare_stage(ref_dir: &Path, name: &str, reference_file: &str, actual: &[f32]) -> StageMetric {
    let reference_path = ref_dir.join(reference_file);
    let expected = read_f32(&reference_path);
    assert_eq!(
        actual.len(),
        expected.len(),
        "stage {name:?} has {} Meganeura values and {} reference values",
        actual.len(),
        expected.len()
    );
    let meganeura_file = format!("meganeura-{reference_file}");
    let meganeura_path = ref_dir.join(&meganeura_file);
    std::fs::write(&meganeura_path, bytemuck::cast_slice(actual))
        .unwrap_or_else(|error| panic!("{}: {error}", meganeura_path.display()));
    StageMetric {
        name: name.to_string(),
        reference_file: reference_file.to_string(),
        meganeura_file,
        elements: actual.len(),
        reference_sha256: common::sha256(&reference_path).expect("hash reference stage"),
        meganeura_sha256: common::sha256(&meganeura_path).expect("hash Meganeura stage"),
        relative_l2: relative_l2(actual, &expected),
        max_absolute_error: actual
            .iter()
            .zip(&expected)
            .map(|(a, b)| (*a - *b).abs())
            .fold(0.0, f32::max),
        cosine: cosine(actual, &expected),
    }
}

fn main() {
    env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("info")).init();

    let mut args = std::env::args().skip(1);
    let ref_dir = PathBuf::from(
        args.next()
            .expect("usage: verify <reference-dir> [model.safetensors]"),
    );
    let weights_path = args
        .next()
        .map(PathBuf::from)
        .unwrap_or_else(|| ref_dir.join("model.safetensors"));

    let pixels = read_f32(&ref_dir.join("pixel_values.bin"));
    let expected = read_f32(&ref_dir.join("features.bin"));
    let reference: ReferenceMetadata = serde_json::from_slice(
        &std::fs::read(ref_dir.join("reference.json")).expect("read reference.json"),
    )
    .expect("parse reference.json");
    let pixel_values_sha256 =
        common::sha256(&ref_dir.join("pixel_values.bin")).expect("hash input");
    let reference_features_sha256 =
        common::sha256(&ref_dir.join("features.bin")).expect("hash reference output");
    let model_sha256 = common::sha256(&weights_path).expect("hash loaded model");
    assert_eq!(
        pixel_values_sha256, reference.pixel_values_sha256,
        "pixel_values.bin hash disagrees with reference.json"
    );
    assert_eq!(
        reference_features_sha256, reference.features_sha256,
        "features.bin hash disagrees with reference.json"
    );
    assert_eq!(
        model_sha256, reference.exported_model_sha256,
        "loaded model hash disagrees with reference.json"
    );

    // Infer the resolution from the pixel tensor rather than assuming it,
    // so the reference script can dump any supported size.
    let side = ((pixels.len() / 3) as f64).sqrt() as usize;
    assert_eq!(
        3 * side * side,
        pixels.len(),
        "pixel tensor is not 3 x N x N"
    );
    assert_eq!(
        side, reference.image_size,
        "pixel tensor resolution disagrees with reference.json"
    );
    let config = Config::vits16()
        .at_resolution(side)
        .with_layers(reference.encoder_layers);
    log::info!(
        "reference: {side}x{side} -> {}x{} grid, {} tokens",
        config.grid(),
        config.grid(),
        config.num_tokens()
    );
    assert_eq!(
        expected.len(),
        config.num_tokens() * config.hidden_size,
        "expected features should be [{}, {}]",
        config.num_tokens(),
        config.hidden_size
    );

    let gpu = dinovision::init_context(None).expect("failed to initialize GPU context");

    let model = meganeura::data::safetensors::SafeTensorsModel::load(weights_path.clone())
        .unwrap_or_else(|e| panic!("{}: {e}", weights_path.display()));
    let patches = dinovision::preprocess::patches_from_pixels_chw(&pixels, &config);
    let mut stages = Vec::new();
    let actual_embeddings = embedding_forward(gpu.clone(), &config, &model, &patches);
    let embedding_metric =
        compare_stage(&ref_dir, "embeddings", "embeddings.bin", &actual_embeddings);
    println!(
        "embedding relative L2 : {:.6}",
        embedding_metric.relative_l2
    );
    stages.push(embedding_metric);
    let actual_first = first_projections(gpu.clone(), &config, &model, &actual_embeddings);
    for (label, actual, file) in [
        ("first norm1", &actual_first[0], "first-norm1.bin"),
        ("first Q", &actual_first[1], "first-q.bin"),
        ("first K", &actual_first[2], "first-k.bin"),
        ("first V", &actual_first[3], "first-v.bin"),
    ] {
        let metric = compare_stage(&ref_dir, label, file, actual);
        println!("{label:<18}: {:.6}", metric.relative_l2);
        stages.push(metric);
    }
    let actual_attention = first_attention(
        gpu.clone(),
        &config,
        &actual_first[1],
        &actual_first[2],
        &actual_first[3],
    );
    for (label, actual, file) in [
        ("first Q RoPE", &actual_attention[0], "first-q-rope.bin"),
        ("first K RoPE", &actual_attention[1], "first-k-rope.bin"),
        (
            "first attention",
            &actual_attention[2],
            "first-attention.bin",
        ),
    ] {
        let metric = compare_stage(&ref_dir, label, file, actual);
        println!("{label:<18}: {:.6}", metric.relative_l2);
        stages.push(metric);
    }
    let actual_remainder = first_remainder(
        gpu.clone(),
        &config,
        &model,
        &actual_embeddings,
        &actual_attention[2],
    );
    for (index, (label, file)) in [
        ("attention projected", "first-attention-projected.bin"),
        ("attention scaled", "first-attention-scaled.bin"),
        ("attention residual", "first-residual.bin"),
        ("first norm2", "first-norm2.bin"),
        ("first MLP up", "first-mlp-up.bin"),
        ("first MLP GELU", "first-mlp-activated.bin"),
        ("first MLP down", "first-mlp-down.bin"),
        ("first MLP scaled", "first-mlp-scaled.bin"),
        ("first output", "first-output.bin"),
        ("first final norm", "first-final-norm.bin"),
    ]
    .into_iter()
    .enumerate()
    {
        let metric = compare_stage(&ref_dir, label, file, &actual_remainder[index]);
        println!("{label:<20}: {:.6}", metric.relative_l2);
        stages.push(metric);
    }

    let (mut session, _) = dinovision::bench::build_encoder_session(gpu, &config, None);
    dinovision::weights::load_encoder(&mut session, &model, &config).expect("load weights");

    session.set_input("patches", &patches);
    session.step();
    session.wait();

    let got = session.read_output(config.num_tokens() * config.hidden_size);

    // Report both absolute error and cosine similarity per token group.
    // Cosine matters more for what we do downstream: PCA colouring and a
    // decoder both care about feature *direction*, and a uniform scale
    // error would still look right while indicating a real bug.
    let mut worst_abs = 0.0f32;
    let mut worst_token = 0usize;
    let mut squared_error = 0.0f64;
    let mut squared_reference = 0.0f64;
    for t in 0..config.num_tokens() {
        for d in 0..config.hidden_size {
            let i = t * config.hidden_size + d;
            let difference = got[i] - expected[i];
            let e = difference.abs();
            squared_error += (difference as f64).powi(2);
            squared_reference += (expected[i] as f64).powi(2);
            if e > worst_abs {
                worst_abs = e;
                worst_token = t;
            }
        }
    }
    let relative_l2 = squared_error.sqrt() / squared_reference.sqrt().max(f64::MIN_POSITIVE);

    let h = config.hidden_size;
    let cls_cos = cosine(&got[0..h], &expected[0..h]);
    let mut worst_patch_cos = 1.0f32;
    for t in config.num_prefix_tokens()..config.num_tokens() {
        let c = cosine(&got[t * h..(t + 1) * h], &expected[t * h..(t + 1) * h]);
        worst_patch_cos = worst_patch_cos.min(c);
    }

    println!("relative L2           : {relative_l2:.6}");
    println!("max |ours - reference| : {worst_abs:.5} (worst at token {worst_token})");
    println!("CLS cosine             : {cls_cos:.6}");
    println!("worst patch cosine     : {worst_patch_cos:.6}");

    // f32 GPU accumulation in a different order than PyTorch's will not
    // reproduce bit-for-bit; 0.999 cosine across every patch is the real
    // signal that the architecture is right.
    let thresholds = Thresholds {
        max_relative_l2: 0.01,
        min_cls_cosine: 0.999,
        min_patch_cosine: 0.999,
    };
    let ok = relative_l2 <= thresholds.max_relative_l2
        && worst_patch_cos > thresholds.min_patch_cosine
        && cls_cos > thresholds.min_cls_cosine;
    let final_metric = compare_stage(&ref_dir, "full encoder output", "features.bin", &got);
    let meganeura_features_sha256 = final_metric.meganeura_sha256.clone();
    stages.push(final_metric);
    let verification = Verification {
        schema_version: 1,
        image_size: config.image_size,
        encoder_layers: config.num_hidden_layers,
        model_sha256,
        source_model_sha256: reference.source_model_sha256,
        pixel_values_sha256,
        reference_features_sha256,
        meganeura_features_sha256,
        elements: got.len(),
        relative_l2,
        max_absolute_error: worst_abs,
        worst_absolute_token: worst_token,
        cls_cosine: cls_cos,
        worst_patch_cosine: worst_patch_cos,
        stages,
        thresholds,
        passed: ok,
    };
    std::fs::write(
        ref_dir.join("verification.json"),
        serde_json::to_vec_pretty(&verification).expect("serialize verification"),
    )
    .expect("write verification record");
    println!("\n{}", if ok { "PASS" } else { "FAIL" });
    if !ok {
        std::process::exit(1);
    }
}