rustyverse/rustytorch
2
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

style: cargo fmt --workspace (whitespace/wrapping only, no semantic change)

Browse Source 
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]>
This commit is contained in:
osobh
2026-08-10 07:09:36 -07:00
co-authored by Claude Sonnet 5
parent ad6405663f
commit 4aaa36a57a
305 changed files with 25537 additions and 18337 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/training/rtx-transformers/src/ssl/jepa_checkpoint.rs
+78 -21
View File
@@ -45,10 +45,17 @@ pub struct JepaModelCheckpoint {
pub enum CheckpointError {
Io(String),
InvalidMagic,
VersionMismatch { found: u32, expected: u32 },
VersionMismatch {
found: u32,
expected: u32,
},
ChecksumMismatch,
MissingField(String),
WrongSize { field: String, expected: usize, found: usize },
WrongSize {
field: String,
expected: usize,
found: usize,
},
}
impl std::fmt::Display for CheckpointError {
@@ -59,10 +66,19 @@ impl std::fmt::Display for CheckpointError {
CheckpointError::VersionMismatch { found, expected } => {
write!(f, "version mismatch: found {found}, expected {expected}")
}
CheckpointError::ChecksumMismatch => write!(f, "checksum mismatch: data may be corrupt"),
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}")
CheckpointError::WrongSize {
field,
expected,
found,
} => {
write!(
f,
"field '{field}' has wrong size: expected {expected}, found {found}"
)
}
}
}
@@ -191,12 +207,16 @@ pub fn deserialize_checkpoint(data: &[u8]) -> Result<JepaModelCheckpoint, Checkp
// Version
let version = read_u32(data, &mut cursor).ok_or(CheckpointError::InvalidMagic)?;
if version != VERSION {
return Err(CheckpointError::VersionMismatch { found: version, expected: 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;
.ok_or_else(|| CheckpointError::Io("truncated at field_count".to_string()))?
as usize;
// Read fields
let mut fields = Vec::with_capacity(field_count);
@@ -205,7 +225,8 @@ pub fn deserialize_checkpoint(data: &[u8]) -> Result<JepaModelCheckpoint, Checkp
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;
.ok_or_else(|| CheckpointError::Io("truncated at name_len".to_string()))?
as usize;
// name bytes
if cursor + name_len > data.len() {
@@ -217,7 +238,8 @@ pub fn deserialize_checkpoint(data: &[u8]) -> Result<JepaModelCheckpoint, Checkp
// value_count
let value_count = read_u32(data, &mut cursor)
.ok_or_else(|| CheckpointError::Io("truncated at value_count".to_string()))? as usize;
.ok_or_else(|| CheckpointError::Io("truncated at value_count".to_string()))?
as usize;
// f32 values
let byte_len = value_count * 4;
@@ -243,7 +265,8 @@ pub fn deserialize_checkpoint(data: &[u8]) -> Result<JepaModelCheckpoint, Checkp
// Step
let step = read_u64(data, &mut cursor)
.ok_or_else(|| CheckpointError::Io("truncated at step".to_string()))? as usize;
.ok_or_else(|| CheckpointError::Io("truncated at step".to_string()))?
as usize;
// Checksum
let checksum_stored = read_u32(data, &mut cursor)
@@ -408,8 +431,8 @@ pub fn apply_fields_to_encoder(
#[cfg(test)]
mod tests {
use super::*;
use super::super::jepa_vit::{CpuViTEncoder, JepaViTConfig};
use super::*;
fn tiny_config() -> JepaViTConfig {
JepaViTConfig {
@@ -558,7 +581,13 @@ mod tests {
bytes[7] = v2[3];
let result = deserialize_checkpoint(&bytes);
assert!(
matches!(result, Err(CheckpointError::VersionMismatch { found: 2, expected: 1 })),
matches!(
result,
Err(CheckpointError::VersionMismatch {
found: 2,
expected: 1
})
),
"expected VersionMismatch{{found:2, expected:1}}, got {result:?}"
);
}
@@ -584,8 +613,14 @@ mod tests {
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] },
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,
@@ -613,7 +648,10 @@ mod tests {
#[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");
assert!(
result.is_err(),
"loading a nonexistent file must return Err"
);
}
// 12. encoder_to_fields produces "patch_embed" field
@@ -653,7 +691,10 @@ mod tests {
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");
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);
@@ -703,16 +744,32 @@ mod tests {
#[test]
fn test_checkpoint_error_display() {
let s = CheckpointError::InvalidMagic.to_string();
assert!(!s.is_empty(), "Display for InvalidMagic should be non-empty");
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 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");
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");
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");