osobhandClaude Sonnet 5
4aaa36a57a
Whole-workspace rustfmt pass picked up while iterating on Mamba GPU backward work. Verified formatting-only via diff sampling; no logic changed. Co-Authored-By: Claude Sonnet 5 <[email protected]>
787 lines
25 KiB
Rust
787 lines
25 KiB
Rust
//! JEPA model checkpoint serialization.
|
||
//!
|
||
//! Binary format: simple, self-describing, no external dependencies.
|
||
//!
|
||
//! File layout:
|
||
//! ```text
|
||
//! [0..4] magic: b"JEPA"
|
||
//! [4..8] version: u32 LE = 1
|
||
//! [8..12] field_count: u32 LE
|
||
//! For each field:
|
||
//! [N..N+4] name_len: u32 LE
|
||
//! [N+4..] name: UTF-8 bytes (name_len bytes)
|
||
//! [..] value_count: u32 LE
|
||
//! [..] values: f32 LE (value_count × 4 bytes)
|
||
//! [end-8] step: u64 LE
|
||
//! [end-4] checksum: u32 LE (sum of all value bytes, wrapping)
|
||
//! ```
|
||
|
||
use super::jepa_vit::CpuViTEncoder;
|
||
|
||
// ============================================================================
|
||
// Public types
|
||
// ============================================================================
|
||
|
||
/// Named weight tensor saved to a JEPA checkpoint.
|
||
#[derive(Debug, Clone)]
|
||
pub struct WeightField {
|
||
pub name: String,
|
||
pub values: Vec<f32>,
|
||
}
|
||
|
||
/// A complete JEPA model checkpoint.
|
||
#[derive(Debug, Clone)]
|
||
pub struct JepaModelCheckpoint {
|
||
pub step: usize,
|
||
pub fields: Vec<WeightField>,
|
||
/// JSON-encoded JepaRunConfig fields for reference.
|
||
pub config_summary: String,
|
||
pub mean_loss: f32,
|
||
pub ema_tau: f32,
|
||
}
|
||
|
||
/// Save/load errors.
|
||
#[derive(Debug)]
|
||
pub enum CheckpointError {
|
||
Io(String),
|
||
InvalidMagic,
|
||
VersionMismatch {
|
||
found: u32,
|
||
expected: u32,
|
||
},
|
||
ChecksumMismatch,
|
||
MissingField(String),
|
||
WrongSize {
|
||
field: String,
|
||
expected: usize,
|
||
found: usize,
|
||
},
|
||
}
|
||
|
||
impl std::fmt::Display for CheckpointError {
|
||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||
match self {
|
||
CheckpointError::Io(msg) => write!(f, "I/O error: {msg}"),
|
||
CheckpointError::InvalidMagic => write!(f, "invalid magic bytes (expected b\"JEPA\")"),
|
||
CheckpointError::VersionMismatch { found, expected } => {
|
||
write!(f, "version mismatch: found {found}, expected {expected}")
|
||
}
|
||
CheckpointError::ChecksumMismatch => {
|
||
write!(f, "checksum mismatch: data may be corrupt")
|
||
}
|
||
CheckpointError::MissingField(name) => write!(f, "missing field: {name}"),
|
||
CheckpointError::WrongSize {
|
||
field,
|
||
expected,
|
||
found,
|
||
} => {
|
||
write!(
|
||
f,
|
||
"field '{field}' has wrong size: expected {expected}, found {found}"
|
||
)
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
// ============================================================================
|
||
// Binary format helpers
|
||
// ============================================================================
|
||
|
||
const MAGIC: &[u8; 4] = b"JEPA";
|
||
const VERSION: u32 = 1;
|
||
|
||
/// Read a u32 from `data` starting at `*cursor` (little-endian). Advances cursor.
|
||
fn read_u32(data: &[u8], cursor: &mut usize) -> Option<u32> {
|
||
if *cursor + 4 > data.len() {
|
||
return None;
|
||
}
|
||
let v = u32::from_le_bytes([
|
||
data[*cursor],
|
||
data[*cursor + 1],
|
||
data[*cursor + 2],
|
||
data[*cursor + 3],
|
||
]);
|
||
*cursor += 4;
|
||
Some(v)
|
||
}
|
||
|
||
/// Read a u64 from `data` starting at `*cursor` (little-endian). Advances cursor.
|
||
fn read_u64(data: &[u8], cursor: &mut usize) -> Option<u64> {
|
||
if *cursor + 8 > data.len() {
|
||
return None;
|
||
}
|
||
let v = u64::from_le_bytes([
|
||
data[*cursor],
|
||
data[*cursor + 1],
|
||
data[*cursor + 2],
|
||
data[*cursor + 3],
|
||
data[*cursor + 4],
|
||
data[*cursor + 5],
|
||
data[*cursor + 6],
|
||
data[*cursor + 7],
|
||
]);
|
||
*cursor += 8;
|
||
Some(v)
|
||
}
|
||
|
||
/// Write a u32 in little-endian to `buf`.
|
||
#[inline]
|
||
fn push_u32(buf: &mut Vec<u8>, v: u32) {
|
||
buf.extend_from_slice(&v.to_le_bytes());
|
||
}
|
||
|
||
/// Write a u64 in little-endian to `buf`.
|
||
#[inline]
|
||
fn push_u64(buf: &mut Vec<u8>, v: u64) {
|
||
buf.extend_from_slice(&v.to_le_bytes());
|
||
}
|
||
|
||
// ============================================================================
|
||
// Serialize / Deserialize
|
||
// ============================================================================
|
||
|
||
/// Serialize a [`JepaModelCheckpoint`] to binary bytes.
|
||
///
|
||
/// Format:
|
||
/// 1. `b"JEPA"` magic (4 bytes)
|
||
/// 2. version = 1 as u32 LE (4 bytes)
|
||
/// 3. field_count as u32 LE (4 bytes)
|
||
/// 4. For each WeightField: name_len u32 LE, name bytes, value_count u32 LE, f32 values
|
||
/// 5. step as u64 LE (8 bytes)
|
||
/// 6. checksum: wrapping sum of all f32 raw bytes, as u32 LE (4 bytes)
|
||
pub fn serialize_checkpoint(ckpt: &JepaModelCheckpoint) -> Vec<u8> {
|
||
let mut buf = Vec::new();
|
||
|
||
// Magic
|
||
buf.extend_from_slice(MAGIC);
|
||
|
||
// Version
|
||
push_u32(&mut buf, VERSION);
|
||
|
||
// Field count
|
||
push_u32(&mut buf, ckpt.fields.len() as u32);
|
||
|
||
// Accumulate checksum across all value bytes
|
||
let mut checksum = 0u32;
|
||
|
||
for field in &ckpt.fields {
|
||
// name_len + name bytes
|
||
let name_bytes = field.name.as_bytes();
|
||
push_u32(&mut buf, name_bytes.len() as u32);
|
||
buf.extend_from_slice(name_bytes);
|
||
|
||
// value_count + f32 values
|
||
push_u32(&mut buf, field.values.len() as u32);
|
||
for &v in &field.values {
|
||
let bytes = v.to_le_bytes();
|
||
// Accumulate each byte into checksum
|
||
for b in bytes {
|
||
checksum = checksum.wrapping_add(b as u32);
|
||
}
|
||
buf.extend_from_slice(&bytes);
|
||
}
|
||
}
|
||
|
||
// Step
|
||
push_u64(&mut buf, ckpt.step as u64);
|
||
|
||
// Checksum
|
||
push_u32(&mut buf, checksum);
|
||
|
||
buf
|
||
}
|
||
|
||
/// Deserialize a [`JepaModelCheckpoint`] from binary bytes.
|
||
///
|
||
/// Returns `Err` on format violations (magic, version, checksum).
|
||
pub fn deserialize_checkpoint(data: &[u8]) -> Result<JepaModelCheckpoint, CheckpointError> {
|
||
let mut cursor = 0usize;
|
||
|
||
// Magic
|
||
if data.len() < 4 || &data[0..4] != MAGIC {
|
||
return Err(CheckpointError::InvalidMagic);
|
||
}
|
||
cursor += 4;
|
||
|
||
// Version
|
||
let version = read_u32(data, &mut cursor).ok_or(CheckpointError::InvalidMagic)?;
|
||
if version != VERSION {
|
||
return Err(CheckpointError::VersionMismatch {
|
||
found: version,
|
||
expected: VERSION,
|
||
});
|
||
}
|
||
|
||
// Field count
|
||
let field_count = read_u32(data, &mut cursor)
|
||
.ok_or_else(|| CheckpointError::Io("truncated at field_count".to_string()))?
|
||
as usize;
|
||
|
||
// Read fields
|
||
let mut fields = Vec::with_capacity(field_count);
|
||
let mut checksum_computed = 0u32;
|
||
|
||
for _ in 0..field_count {
|
||
// name_len
|
||
let name_len = read_u32(data, &mut cursor)
|
||
.ok_or_else(|| CheckpointError::Io("truncated at name_len".to_string()))?
|
||
as usize;
|
||
|
||
// name bytes
|
||
if cursor + name_len > data.len() {
|
||
return Err(CheckpointError::Io("truncated at name bytes".to_string()));
|
||
}
|
||
let name = String::from_utf8(data[cursor..cursor + name_len].to_vec())
|
||
.map_err(|e| CheckpointError::Io(format!("invalid UTF-8 in field name: {e}")))?;
|
||
cursor += name_len;
|
||
|
||
// value_count
|
||
let value_count = read_u32(data, &mut cursor)
|
||
.ok_or_else(|| CheckpointError::Io("truncated at value_count".to_string()))?
|
||
as usize;
|
||
|
||
// f32 values
|
||
let byte_len = value_count * 4;
|
||
if cursor + byte_len > data.len() {
|
||
return Err(CheckpointError::Io(format!(
|
||
"truncated reading values for field '{name}': need {byte_len} bytes at offset {cursor}, have {}",
|
||
data.len()
|
||
)));
|
||
}
|
||
let mut values = Vec::with_capacity(value_count);
|
||
for i in 0..value_count {
|
||
let off = cursor + i * 4;
|
||
let bytes = [data[off], data[off + 1], data[off + 2], data[off + 3]];
|
||
for b in bytes {
|
||
checksum_computed = checksum_computed.wrapping_add(b as u32);
|
||
}
|
||
values.push(f32::from_le_bytes(bytes));
|
||
}
|
||
cursor += byte_len;
|
||
|
||
fields.push(WeightField { name, values });
|
||
}
|
||
|
||
// Step
|
||
let step = read_u64(data, &mut cursor)
|
||
.ok_or_else(|| CheckpointError::Io("truncated at step".to_string()))?
|
||
as usize;
|
||
|
||
// Checksum
|
||
let checksum_stored = read_u32(data, &mut cursor)
|
||
.ok_or_else(|| CheckpointError::Io("truncated at checksum".to_string()))?;
|
||
|
||
if checksum_stored != checksum_computed {
|
||
return Err(CheckpointError::ChecksumMismatch);
|
||
}
|
||
|
||
Ok(JepaModelCheckpoint {
|
||
step,
|
||
fields,
|
||
config_summary: String::new(),
|
||
mean_loss: 0.0,
|
||
ema_tau: 0.0,
|
||
})
|
||
}
|
||
|
||
// ============================================================================
|
||
// File I/O
|
||
// ============================================================================
|
||
|
||
/// Write a checkpoint to disk at `path`.
|
||
pub fn save_checkpoint(ckpt: &JepaModelCheckpoint, path: &str) -> Result<(), CheckpointError> {
|
||
let bytes = serialize_checkpoint(ckpt);
|
||
std::fs::write(path, &bytes)
|
||
.map_err(|e| CheckpointError::Io(format!("failed to write '{path}': {e}")))?;
|
||
Ok(())
|
||
}
|
||
|
||
/// Load a checkpoint from disk.
|
||
pub fn load_checkpoint(path: &str) -> Result<JepaModelCheckpoint, CheckpointError> {
|
||
let data = std::fs::read(path)
|
||
.map_err(|e| CheckpointError::Io(format!("failed to read '{path}': {e}")))?;
|
||
deserialize_checkpoint(&data)
|
||
}
|
||
|
||
// ============================================================================
|
||
// Encoder <-> WeightField conversion
|
||
// ============================================================================
|
||
|
||
/// Extract weights from a [`CpuViTEncoder`] into a list of [`WeightField`]s.
|
||
///
|
||
/// Field names:
|
||
/// - `"patch_embed"` — `[num_patches * embed_dim]` sinusoidal PE + projection
|
||
/// - `"proj_w"` — `[embed_dim * embed_dim]` patch projection weights
|
||
/// - `"proj_b"` — `[embed_dim]` patch projection bias
|
||
/// - `"block_{i}_qkv_w"` — `[embed_dim * 3 * embed_dim]`
|
||
/// - `"block_{i}_qkv_b"` — `[3 * embed_dim]`
|
||
/// - `"block_{i}_out_w"` — `[embed_dim * embed_dim]`
|
||
/// - `"block_{i}_out_b"` — `[embed_dim]`
|
||
/// - `"block_{i}_ffn1_w"` — `[embed_dim * ffn_dim]`
|
||
/// - `"block_{i}_ffn1_b"` — `[ffn_dim]`
|
||
/// - `"block_{i}_ffn2_w"` — `[ffn_dim * embed_dim]`
|
||
/// - `"block_{i}_ffn2_b"` — `[embed_dim]`
|
||
pub fn encoder_to_fields(encoder: &CpuViTEncoder) -> Vec<WeightField> {
|
||
let mut fields = Vec::new();
|
||
|
||
fields.push(WeightField {
|
||
name: "patch_embed".to_string(),
|
||
values: encoder.patch_embed.clone(),
|
||
});
|
||
fields.push(WeightField {
|
||
name: "proj_w".to_string(),
|
||
values: encoder.proj_w.clone(),
|
||
});
|
||
fields.push(WeightField {
|
||
name: "proj_b".to_string(),
|
||
values: encoder.proj_b.clone(),
|
||
});
|
||
|
||
for (i, block) in encoder.blocks.iter().enumerate() {
|
||
fields.push(WeightField {
|
||
name: format!("block_{i}_qkv_w"),
|
||
values: block.qkv_w.clone(),
|
||
});
|
||
fields.push(WeightField {
|
||
name: format!("block_{i}_qkv_b"),
|
||
values: block.qkv_b.clone(),
|
||
});
|
||
fields.push(WeightField {
|
||
name: format!("block_{i}_out_w"),
|
||
values: block.out_w.clone(),
|
||
});
|
||
fields.push(WeightField {
|
||
name: format!("block_{i}_out_b"),
|
||
values: block.out_b.clone(),
|
||
});
|
||
fields.push(WeightField {
|
||
name: format!("block_{i}_ffn1_w"),
|
||
values: block.ffn1_w.clone(),
|
||
});
|
||
fields.push(WeightField {
|
||
name: format!("block_{i}_ffn1_b"),
|
||
values: block.ffn1_b.clone(),
|
||
});
|
||
fields.push(WeightField {
|
||
name: format!("block_{i}_ffn2_w"),
|
||
values: block.ffn2_w.clone(),
|
||
});
|
||
fields.push(WeightField {
|
||
name: format!("block_{i}_ffn2_b"),
|
||
values: block.ffn2_b.clone(),
|
||
});
|
||
}
|
||
|
||
fields
|
||
}
|
||
|
||
/// Apply [`WeightField`]s back into a [`CpuViTEncoder`] in-place.
|
||
///
|
||
/// Returns `Err(MissingField)` if a required field is absent,
|
||
/// or `Err(WrongSize)` if a field has incorrect length.
|
||
pub fn apply_fields_to_encoder(
|
||
encoder: &mut CpuViTEncoder,
|
||
fields: &[WeightField],
|
||
) -> Result<(), CheckpointError> {
|
||
// Build a lookup map from field name to values slice
|
||
let map: std::collections::HashMap<&str, &[f32]> = fields
|
||
.iter()
|
||
.map(|f| (f.name.as_str(), f.values.as_slice()))
|
||
.collect();
|
||
|
||
macro_rules! apply_field {
|
||
($name:expr, $target:expr) => {{
|
||
let name: &str = $name;
|
||
let vals = map
|
||
.get(name)
|
||
.ok_or_else(|| CheckpointError::MissingField(name.to_string()))?;
|
||
if vals.len() != $target.len() {
|
||
return Err(CheckpointError::WrongSize {
|
||
field: name.to_string(),
|
||
expected: $target.len(),
|
||
found: vals.len(),
|
||
});
|
||
}
|
||
$target.copy_from_slice(vals);
|
||
}};
|
||
}
|
||
|
||
apply_field!("patch_embed", encoder.patch_embed);
|
||
apply_field!("proj_w", encoder.proj_w);
|
||
apply_field!("proj_b", encoder.proj_b);
|
||
|
||
for (i, block) in encoder.blocks.iter_mut().enumerate() {
|
||
apply_field!(&format!("block_{i}_qkv_w"), block.qkv_w);
|
||
apply_field!(&format!("block_{i}_qkv_b"), block.qkv_b);
|
||
apply_field!(&format!("block_{i}_out_w"), block.out_w);
|
||
apply_field!(&format!("block_{i}_out_b"), block.out_b);
|
||
apply_field!(&format!("block_{i}_ffn1_w"), block.ffn1_w);
|
||
apply_field!(&format!("block_{i}_ffn1_b"), block.ffn1_b);
|
||
apply_field!(&format!("block_{i}_ffn2_w"), block.ffn2_w);
|
||
apply_field!(&format!("block_{i}_ffn2_b"), block.ffn2_b);
|
||
}
|
||
|
||
Ok(())
|
||
}
|
||
|
||
// ============================================================================
|
||
// Tests
|
||
// ============================================================================
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::super::jepa_vit::{CpuViTEncoder, JepaViTConfig};
|
||
use super::*;
|
||
|
||
fn tiny_config() -> JepaViTConfig {
|
||
JepaViTConfig {
|
||
embed_dim: 32,
|
||
depth: 2,
|
||
num_heads: 4,
|
||
mlp_ratio: 2.0,
|
||
patch_size: 16,
|
||
image_size: 64,
|
||
}
|
||
}
|
||
|
||
fn empty_ckpt(step: usize) -> JepaModelCheckpoint {
|
||
JepaModelCheckpoint {
|
||
step,
|
||
fields: Vec::new(),
|
||
config_summary: String::new(),
|
||
mean_loss: 0.0,
|
||
ema_tau: 0.0,
|
||
}
|
||
}
|
||
|
||
fn single_field_ckpt(step: usize) -> JepaModelCheckpoint {
|
||
JepaModelCheckpoint {
|
||
step,
|
||
fields: vec![WeightField {
|
||
name: "test".to_string(),
|
||
values: vec![1.0, 2.0, 3.0],
|
||
}],
|
||
config_summary: String::new(),
|
||
mean_loss: 0.0,
|
||
ema_tau: 0.0,
|
||
}
|
||
}
|
||
|
||
// 1. Magic bytes
|
||
#[test]
|
||
fn test_magic_bytes() {
|
||
let ckpt = empty_ckpt(0);
|
||
let bytes = serialize_checkpoint(&ckpt);
|
||
assert_eq!(&bytes[0..4], b"JEPA", "first 4 bytes must be b\"JEPA\"");
|
||
}
|
||
|
||
// 2. Version
|
||
#[test]
|
||
fn test_version() {
|
||
let ckpt = empty_ckpt(0);
|
||
let bytes = serialize_checkpoint(&ckpt);
|
||
let v = u32::from_le_bytes([bytes[4], bytes[5], bytes[6], bytes[7]]);
|
||
assert_eq!(v, 1, "version field must be 1");
|
||
}
|
||
|
||
// 3. Round-trip: empty fields
|
||
#[test]
|
||
fn test_round_trip_empty_fields() {
|
||
let ckpt = empty_ckpt(42);
|
||
let bytes = serialize_checkpoint(&ckpt);
|
||
let loaded = deserialize_checkpoint(&bytes).expect("deserialize should succeed");
|
||
assert_eq!(loaded.step, 42);
|
||
assert_eq!(loaded.fields.len(), 0);
|
||
}
|
||
|
||
// 4. Round-trip: single field
|
||
#[test]
|
||
fn test_round_trip_single_field() {
|
||
let ckpt = single_field_ckpt(7);
|
||
let bytes = serialize_checkpoint(&ckpt);
|
||
let loaded = deserialize_checkpoint(&bytes).expect("deserialize should succeed");
|
||
assert_eq!(loaded.step, 7);
|
||
assert_eq!(loaded.fields.len(), 1);
|
||
assert_eq!(loaded.fields[0].name, "test");
|
||
assert_eq!(loaded.fields[0].values, vec![1.0f32, 2.0, 3.0]);
|
||
}
|
||
|
||
// 5. Round-trip: multiple fields
|
||
#[test]
|
||
fn test_round_trip_multiple_fields() {
|
||
let fields: Vec<WeightField> = (0..5)
|
||
.map(|i| WeightField {
|
||
name: format!("field_{i}"),
|
||
values: (0..i * 10 + 1).map(|j| j as f32 * 0.1).collect(),
|
||
})
|
||
.collect();
|
||
let ckpt = JepaModelCheckpoint {
|
||
step: 100,
|
||
fields,
|
||
config_summary: "test".to_string(),
|
||
mean_loss: 1.23,
|
||
ema_tau: 0.996,
|
||
};
|
||
let bytes = serialize_checkpoint(&ckpt);
|
||
let loaded = deserialize_checkpoint(&bytes).expect("deserialize should succeed");
|
||
assert_eq!(loaded.step, 100);
|
||
assert_eq!(loaded.fields.len(), 5);
|
||
for i in 0..5 {
|
||
assert_eq!(loaded.fields[i].name, format!("field_{i}"));
|
||
assert_eq!(loaded.fields[i].values.len(), i * 10 + 1);
|
||
}
|
||
}
|
||
|
||
// 6. Round-trip: large field
|
||
#[test]
|
||
fn test_round_trip_large_field() {
|
||
let n = 50_000;
|
||
let ckpt = JepaModelCheckpoint {
|
||
step: 999,
|
||
fields: vec![WeightField {
|
||
name: "large".to_string(),
|
||
values: (0..n).map(|i| i as f32 * 1e-5).collect(),
|
||
}],
|
||
config_summary: String::new(),
|
||
mean_loss: 0.0,
|
||
ema_tau: 0.0,
|
||
};
|
||
let bytes = serialize_checkpoint(&ckpt);
|
||
let loaded = deserialize_checkpoint(&bytes).expect("large field should round-trip");
|
||
assert_eq!(loaded.fields[0].values.len(), n);
|
||
// Spot-check a few values
|
||
assert!((loaded.fields[0].values[0] - 0.0).abs() < 1e-10);
|
||
assert!((loaded.fields[0].values[1] - 1e-5).abs() < 1e-12);
|
||
}
|
||
|
||
// 7. Invalid magic
|
||
#[test]
|
||
fn test_invalid_magic() {
|
||
let ckpt = empty_ckpt(0);
|
||
let mut bytes = serialize_checkpoint(&ckpt);
|
||
bytes[0] = 0xFF; // corrupt first magic byte
|
||
let result = deserialize_checkpoint(&bytes);
|
||
assert!(
|
||
matches!(result, Err(CheckpointError::InvalidMagic)),
|
||
"expected InvalidMagic, got {result:?}"
|
||
);
|
||
}
|
||
|
||
// 8. Version mismatch
|
||
#[test]
|
||
fn test_version_mismatch() {
|
||
let ckpt = empty_ckpt(0);
|
||
let mut bytes = serialize_checkpoint(&ckpt);
|
||
// Overwrite version bytes (offset 4..8) with version=2
|
||
let v2 = 2u32.to_le_bytes();
|
||
bytes[4] = v2[0];
|
||
bytes[5] = v2[1];
|
||
bytes[6] = v2[2];
|
||
bytes[7] = v2[3];
|
||
let result = deserialize_checkpoint(&bytes);
|
||
assert!(
|
||
matches!(
|
||
result,
|
||
Err(CheckpointError::VersionMismatch {
|
||
found: 2,
|
||
expected: 1
|
||
})
|
||
),
|
||
"expected VersionMismatch{{found:2, expected:1}}, got {result:?}"
|
||
);
|
||
}
|
||
|
||
// 9. Checksum corruption
|
||
#[test]
|
||
fn test_checksum_corruption() {
|
||
let ckpt = single_field_ckpt(1);
|
||
let mut bytes = serialize_checkpoint(&ckpt);
|
||
// Flip the very last byte (checksum)
|
||
let last = bytes.len() - 1;
|
||
bytes[last] ^= 0xFF;
|
||
let result = deserialize_checkpoint(&bytes);
|
||
assert!(
|
||
matches!(result, Err(CheckpointError::ChecksumMismatch)),
|
||
"expected ChecksumMismatch, got {result:?}"
|
||
);
|
||
}
|
||
|
||
// 10. Save / load round-trip via temp file
|
||
#[test]
|
||
fn test_save_load_roundtrip() {
|
||
let ckpt = JepaModelCheckpoint {
|
||
step: 12345,
|
||
fields: vec![
|
||
WeightField {
|
||
name: "w1".to_string(),
|
||
values: vec![1.0, 2.0, 3.0],
|
||
},
|
||
WeightField {
|
||
name: "w2".to_string(),
|
||
values: vec![4.0, 5.0],
|
||
},
|
||
],
|
||
config_summary: "{}".to_string(),
|
||
mean_loss: 0.42,
|
||
ema_tau: 0.999,
|
||
};
|
||
|
||
let tmp = std::env::temp_dir().join("jepa_test_save_load.jepa");
|
||
let path = tmp.to_str().expect("temp path is valid UTF-8");
|
||
|
||
save_checkpoint(&ckpt, path).expect("save_checkpoint should succeed");
|
||
let loaded = load_checkpoint(path).expect("load_checkpoint should succeed");
|
||
|
||
assert_eq!(loaded.step, 12345);
|
||
assert_eq!(loaded.fields.len(), 2);
|
||
assert_eq!(loaded.fields[0].name, "w1");
|
||
assert_eq!(loaded.fields[0].values, vec![1.0f32, 2.0, 3.0]);
|
||
assert_eq!(loaded.fields[1].name, "w2");
|
||
assert_eq!(loaded.fields[1].values, vec![4.0f32, 5.0]);
|
||
|
||
// Clean up
|
||
let _ = std::fs::remove_file(path);
|
||
}
|
||
|
||
// 11. Load nonexistent file → Err
|
||
#[test]
|
||
fn test_load_nonexistent_file() {
|
||
let result = load_checkpoint("nonexistent_jepa_file_xyz_abc.jepa");
|
||
assert!(
|
||
result.is_err(),
|
||
"loading a nonexistent file must return Err"
|
||
);
|
||
}
|
||
|
||
// 12. encoder_to_fields produces "patch_embed" field
|
||
#[test]
|
||
fn test_encoder_to_fields_names() {
|
||
let enc = CpuViTEncoder::new(tiny_config());
|
||
let fields = encoder_to_fields(&enc);
|
||
let names: Vec<&str> = fields.iter().map(|f| f.name.as_str()).collect();
|
||
assert!(
|
||
names.contains(&"patch_embed"),
|
||
"fields must contain 'patch_embed'; got: {names:?}"
|
||
);
|
||
}
|
||
|
||
// 13. encoder_to_fields count (tiny has depth=2 → 3 + 8*2 = 19 fields)
|
||
#[test]
|
||
fn test_encoder_to_fields_count() {
|
||
let enc = CpuViTEncoder::new(tiny_config());
|
||
let fields = encoder_to_fields(&enc);
|
||
// 3 base fields (patch_embed, proj_w, proj_b) + 8 per block × depth
|
||
let expected = 3 + 8 * tiny_config().depth;
|
||
assert_eq!(
|
||
fields.len(),
|
||
expected,
|
||
"expected {expected} fields for tiny config (depth=2), got {}",
|
||
fields.len()
|
||
);
|
||
}
|
||
|
||
// 14. apply_fields_to_encoder round-trip: extract from A, apply to B → same weights
|
||
#[test]
|
||
fn test_apply_fields_roundtrip() {
|
||
let enc_a = CpuViTEncoder::new(tiny_config());
|
||
let mut enc_b = CpuViTEncoder::new(tiny_config());
|
||
|
||
// Modify enc_b's patch_embed so they start different
|
||
for v in enc_b.patch_embed.iter_mut() {
|
||
*v = 0.0;
|
||
}
|
||
assert_ne!(
|
||
enc_a.patch_embed, enc_b.patch_embed,
|
||
"encoders should start different"
|
||
);
|
||
|
||
// Extract from A, apply to B
|
||
let fields = encoder_to_fields(&enc_a);
|
||
apply_fields_to_encoder(&mut enc_b, &fields).expect("apply_fields should succeed");
|
||
|
||
// Now B's patch_embed should match A's
|
||
assert_eq!(
|
||
enc_a.patch_embed, enc_b.patch_embed,
|
||
"patch_embed should match after apply_fields"
|
||
);
|
||
}
|
||
|
||
// 15. apply_fields_to_encoder: missing required field → MissingField error
|
||
#[test]
|
||
fn test_apply_fields_missing() {
|
||
let mut enc = CpuViTEncoder::new(tiny_config());
|
||
// Provide only a subset of fields (missing patch_embed)
|
||
let fields: Vec<WeightField> = Vec::new();
|
||
let result = apply_fields_to_encoder(&mut enc, &fields);
|
||
assert!(
|
||
matches!(result, Err(CheckpointError::MissingField(ref name)) if name == "patch_embed"),
|
||
"expected MissingField(\"patch_embed\"), got {result:?}"
|
||
);
|
||
}
|
||
|
||
// 16. apply_fields_to_encoder: wrong size → WrongSize error
|
||
#[test]
|
||
fn test_apply_fields_wrong_size() {
|
||
let mut enc = CpuViTEncoder::new(tiny_config());
|
||
// Build a mostly-correct field list but wrong size for patch_embed
|
||
let mut fields = encoder_to_fields(&enc);
|
||
// Corrupt patch_embed size
|
||
for f in fields.iter_mut() {
|
||
if f.name == "patch_embed" {
|
||
f.values = vec![1.0, 2.0]; // wrong length
|
||
break;
|
||
}
|
||
}
|
||
let result = apply_fields_to_encoder(&mut enc, &fields);
|
||
assert!(
|
||
matches!(result, Err(CheckpointError::WrongSize { ref field, .. }) if field == "patch_embed"),
|
||
"expected WrongSize for patch_embed, got {result:?}"
|
||
);
|
||
}
|
||
|
||
// 17. CheckpointError::Display is non-empty
|
||
#[test]
|
||
fn test_checkpoint_error_display() {
|
||
let s = CheckpointError::InvalidMagic.to_string();
|
||
assert!(
|
||
!s.is_empty(),
|
||
"Display for InvalidMagic should be non-empty"
|
||
);
|
||
|
||
let s2 = CheckpointError::VersionMismatch {
|
||
found: 2,
|
||
expected: 1,
|
||
}
|
||
.to_string();
|
||
assert!(
|
||
!s2.is_empty(),
|
||
"Display for VersionMismatch should be non-empty"
|
||
);
|
||
|
||
let s3 = CheckpointError::ChecksumMismatch.to_string();
|
||
assert!(
|
||
!s3.is_empty(),
|
||
"Display for ChecksumMismatch should be non-empty"
|
||
);
|
||
|
||
let s4 = CheckpointError::MissingField("foo".to_string()).to_string();
|
||
assert!(
|
||
!s4.is_empty(),
|
||
"Display for MissingField should be non-empty"
|
||
);
|
||
|
||
let s5 = CheckpointError::Io("disk full".to_string()).to_string();
|
||
assert!(!s5.is_empty(), "Display for Io should be non-empty");
|
||
}
|
||
|
||
// 18. Step is preserved exactly
|
||
#[test]
|
||
fn test_step_preserved() {
|
||
let ckpt = empty_ckpt(99999);
|
||
let bytes = serialize_checkpoint(&ckpt);
|
||
let loaded = deserialize_checkpoint(&bytes).expect("deserialize should succeed");
|
||
assert_eq!(loaded.step, 99999, "step must be preserved exactly");
|
||
}
|
||
}
|