fn create_binary_tensor_entry(
name: &str,
dtype: u8,
shape: &[u64],
offset: u64,
size: u64,
) -> Vec<u8> {
let mut data = Vec::new();
data.extend_from_slice(&(name.len() as u16).to_le_bytes());
data.extend_from_slice(name.as_bytes());
data.push(dtype);
data.push(shape.len() as u8);
for &dim in shape {
data.extend_from_slice(&dim.to_le_bytes());
}
data.extend_from_slice(&offset.to_le_bytes());
data.extend_from_slice(&size.to_le_bytes());
data
}
fn create_test_apr_model() -> Vec<u8> {
let metadata = r#"{"architecture":"test","vocab_size":100,"hidden_size":64}"#;
let metadata_bytes = metadata.as_bytes();
let metadata_padded_size = metadata_bytes.len().div_ceil(64) * 64;
let tensor_entry = create_binary_tensor_entry("test.weight", 0, &[4, 4], 0, 64);
let tensor_index_offset = HEADER_SIZE as u64 + metadata_padded_size as u64;
let data_offset = tensor_index_offset + tensor_entry.len() as u64;
let data_size = 64usize;
let total_size = data_offset as usize + data_size;
let mut data = vec![0u8; total_size];
data[0..4].copy_from_slice(&MAGIC);
data[4] = 2; data[5] = 0; data[6..8].copy_from_slice(&0u16.to_le_bytes()); data[8..12].copy_from_slice(&1u32.to_le_bytes()); data[12..20].copy_from_slice(&(HEADER_SIZE as u64).to_le_bytes()); data[20..24].copy_from_slice(&(metadata_bytes.len() as u32).to_le_bytes()); data[24..32].copy_from_slice(&tensor_index_offset.to_le_bytes()); data[32..40].copy_from_slice(&data_offset.to_le_bytes());
data[HEADER_SIZE..HEADER_SIZE + metadata_bytes.len()].copy_from_slice(metadata_bytes);
let idx_start = tensor_index_offset as usize;
data[idx_start..idx_start + tensor_entry.len()].copy_from_slice(&tensor_entry);
let data_start = data_offset as usize;
for i in 0..16 {
let val = i as f32;
data[data_start + i * 4..data_start + i * 4 + 4].copy_from_slice(&val.to_le_bytes());
}
data
}
fn create_test_apr_model_with_dtype(dtype: u8, data_bytes: &[u8]) -> Vec<u8> {
let metadata = r#"{"architecture":"test"}"#;
let metadata_bytes = metadata.as_bytes();
let metadata_padded_size = metadata_bytes.len().div_ceil(64) * 64;
let tensor_entry =
create_binary_tensor_entry("typed.weight", dtype, &[4], 0, data_bytes.len() as u64);
let tensor_index_offset = HEADER_SIZE as u64 + metadata_padded_size as u64;
let data_offset = tensor_index_offset + tensor_entry.len() as u64;
let total_size = data_offset as usize + data_bytes.len();
let mut data = vec![0u8; total_size];
data[0..4].copy_from_slice(&MAGIC);
data[4] = 2; data[5] = 0; data[6..8].copy_from_slice(&0u16.to_le_bytes()); data[8..12].copy_from_slice(&1u32.to_le_bytes()); data[12..20].copy_from_slice(&(HEADER_SIZE as u64).to_le_bytes()); data[20..24].copy_from_slice(&(metadata_bytes.len() as u32).to_le_bytes()); data[24..32].copy_from_slice(&tensor_index_offset.to_le_bytes()); data[32..40].copy_from_slice(&data_offset.to_le_bytes());
data[HEADER_SIZE..HEADER_SIZE + metadata_bytes.len()].copy_from_slice(metadata_bytes);
let idx_start = tensor_index_offset as usize;
data[idx_start..idx_start + tensor_entry.len()].copy_from_slice(&tensor_entry);
let data_start = data_offset as usize;
data[data_start..data_start + data_bytes.len()].copy_from_slice(data_bytes);
data
}
#[test]
fn test_rms_norm_basic() {
let x = vec![1.0, 2.0, 3.0, 4.0];
let weight = vec![1.0, 1.0, 1.0, 1.0];
let eps = 1e-6;
let result = crate::apr::rms_norm(&x, &weight, eps);
assert_eq!(result.len(), 4);
}
#[test]
fn test_rms_norm_zeros() {
let x = vec![0.0, 0.0, 0.0, 0.0];
let weight = vec![1.0, 1.0, 1.0, 1.0];
let eps = 1e-6;
let result = crate::apr::rms_norm(&x, &weight, eps);
assert_eq!(result.len(), 4);
for &v in &result {
assert!(v.abs() < 1e-3);
}
}
#[test]
fn test_rms_norm_with_weight() {
let x = vec![1.0, 1.0, 1.0, 1.0];
let weight = vec![2.0, 2.0, 2.0, 2.0];
let eps = 1e-6;
let result = crate::apr::rms_norm(&x, &weight, eps);
assert_eq!(result.len(), 4);
assert!(result[0] > 1.0);
}
#[test]
fn test_matmul_identity() {
let x = vec![1.0, 2.0, 1.0, 2.0];
let w = vec![1.0, 0.0, 0.0, 1.0]; let result = crate::apr::matmul(&x, &w, 2, 2, 2);
assert_eq!(result.len(), 4);
assert!((result[0] - 1.0).abs() < 1e-6);
assert!((result[1] - 2.0).abs() < 1e-6);
}
#[test]
fn test_matmul_simple() {
let x = vec![1.0, 2.0];
let w = vec![1.0, 1.0];
let result = crate::apr::matmul(&x, &w, 1, 2, 1);
assert_eq!(result.len(), 1);
assert!((result[0] - 3.0).abs() < 1e-6);
}
#[test]
fn test_simd_dot_basic() {
let a = vec![1.0, 2.0, 3.0, 4.0];
let b = vec![1.0, 1.0, 1.0, 1.0];
let result = crate::apr::simd_dot(&a, &b);
assert!((result - 10.0).abs() < 1e-6);
}
#[test]
fn test_simd_dot_zeros() {
let a = vec![1.0, 2.0, 3.0, 4.0];
let b = vec![0.0, 0.0, 0.0, 0.0];
let result = crate::apr::simd_dot(&a, &b);
assert!((result - 0.0).abs() < 1e-6);
}
#[test]
fn test_simd_dot_large() {
let a: Vec<f32> = (0..256).map(|i| i as f32).collect();
let b: Vec<f32> = vec![1.0; 256];
let result = crate::apr::simd_dot(&a, &b);
assert!((result - 32640.0).abs() < 1e-3);
}
#[test]
fn test_dequantize_f16_basic() {
let bytes = vec![0x00, 0x3C, 0x00, 0x3C];
let result = crate::apr::dequantize_f16(&bytes, 2);
assert_eq!(result.len(), 2);
assert!((result[0] - 1.0).abs() < 0.01);
assert!((result[1] - 1.0).abs() < 0.01);
}
#[test]
fn test_dequantize_f16_truncated() {
let bytes = vec![0x00]; let result = crate::apr::dequantize_f16(&bytes, 1);
assert_eq!(result.len(), 0);
}
#[test]
fn test_dequantize_q8_0_zero_scale() {
let mut bytes = vec![0u8; 34];
bytes[0] = 0x00;
bytes[1] = 0x00;
let result = crate::apr::dequantize_q8_0(&bytes, 32);
assert_eq!(result.len(), 32);
for &v in &result {
assert!((v - 0.0).abs() < 1e-6);
}
}
#[test]
fn test_tensor_entry_byte_size_f32() {
let entry = TensorEntry {
name: "test".to_string(),
dtype: "F32".to_string(),
shape: vec![2, 3],
offset: 0,
size: 24, };
assert_eq!(entry.element_count(), 6);
}
#[test]
fn test_tensor_entry_byte_size_f16() {
let entry = TensorEntry {
name: "test".to_string(),
dtype: "F16".to_string(),
shape: vec![4, 4],
offset: 0,
size: 32, };
assert_eq!(entry.element_count(), 16);
}
#[test]
fn test_apr_metadata_is_transformer_missing_heads() {
let meta = AprMetadata {
hidden_size: Some(256),
num_layers: Some(4),
num_heads: None,
vocab_size: Some(32000),
..Default::default()
};
assert!(!meta.is_transformer());
}
#[test]
fn test_apr_metadata_is_transformer_missing_vocab() {
let meta = AprMetadata {
hidden_size: Some(256),
num_layers: Some(4),
num_heads: Some(8),
vocab_size: None,
..Default::default()
};
assert!(!meta.is_transformer());
}
#[test]
fn test_simple_attention_single_head() {
let q = vec![1.0, 0.0];
let k = vec![1.0, 0.0];
let v = vec![1.0, 1.0];
let result = crate::apr::simple_attention(&q, &k, &v, 1, 1, 1, 2);
assert_eq!(result.len(), 2);
}
#[test]
fn test_apr_v2_model_from_bytes_truncated() {
let data = vec![0u8; 10];
let result = AprV2Model::from_bytes(data);
assert!(result.is_err());
}
#[test]
fn test_apr_v2_model_from_bytes_invalid_magic() {
let mut data = vec![0u8; 100];
data[0..4].copy_from_slice(b"GGUF");
let result = AprV2Model::from_bytes(data);
assert!(result.is_err());
}
#[test]
fn test_apr_v2_model_get_tensor_bytes_existing() {
let data = create_test_apr_model();
let model = AprV2Model::from_bytes(data).expect("APR operation failed");
let result = model.get_tensor_bytes("test.weight");
assert!(result.is_ok());
}
#[test]
fn test_apr_v2_model_get_tensor_f32_existing() {
let data = create_test_apr_model();
let model = AprV2Model::from_bytes(data).expect("APR operation failed");
let result = model.get_tensor_f32("test.weight");
assert!(result.is_ok());
let tensor = result.expect("APR operation failed");
assert!(!tensor.is_empty());
}
#[test]
fn test_apr_header_magic_validation() {
let mut data = vec![0u8; 64];
data[0..4].copy_from_slice(b"XXXX"); let result = AprHeader::from_bytes(&data);
assert!(result.is_err());
}
#[test]
fn test_apr_header_version_tuple() {
let mut data = vec![0u8; 64];
data[0..4].copy_from_slice(&MAGIC);
data[4] = 2; data[5] = 1; let result = AprHeader::from_bytes(&data);
assert!(result.is_ok());
let header = result.expect("APR operation failed");
assert_eq!(header.version, (2, 1));
}
#[test]
fn test_apr_flags_from_default() {
let flags = AprFlags::default();
assert!(!flags.is_compressed());
assert!(!flags.is_encrypted());
}
#[test]
fn test_tensor_entry_from_binary_large_shape() {
let entry = create_binary_tensor_entry("big_tensor", 0, &[10, 20, 30, 40], 0, 240000);
let (parsed, consumed) = TensorEntry::from_binary(&entry).expect("APR operation failed");
assert_eq!(parsed.name, "big_tensor");
assert_eq!(parsed.shape, vec![10, 20, 30, 40]);
assert_eq!(parsed.element_count(), 10 * 20 * 30 * 40);
assert!(consumed > 0);
}
#[test]
fn test_tensor_entry_from_binary_zero_dim() {
let entry = create_binary_tensor_entry("scalar", 0, &[], 0, 4);
let (parsed, _) = TensorEntry::from_binary(&entry).expect("APR operation failed");
assert_eq!(parsed.name, "scalar");
assert!(parsed.shape.is_empty());
assert_eq!(parsed.element_count(), 1); }
#[test]
fn test_dequantize_f16_multiple_values() {
let bytes = vec![
0x00, 0x3C, 0x00, 0x40, 0x00, 0x3C, ];
let result = crate::apr::dequantize_f16(&bytes, 3);
assert_eq!(result.len(), 3);
assert!((result[0] - 1.0).abs() < 0.1);
assert!((result[1] - 2.0).abs() < 0.1);
}