mod ops;
mod sample_type;
mod tensor;
pub use sample_type::{Sample, SampleMetadata};
pub use tensor::{DType, Tensor};
#[cfg(test)]
mod tests {
use super::*;
use crate::error::Error;
use std::sync::{Arc, OnceLock};
#[test]
fn sample_builder() {
let sample = Sample::new()
.with("x", Tensor::f32(&[1.0, 2.0, 3.0], vec![3]))
.with("y", Tensor::i64(&[42], vec![1]));
assert_eq!(sample.len(), 2);
assert_eq!(sample.get("x").unwrap().shape(), &[3]);
assert_eq!(sample.get("y").unwrap().try_as_i64().unwrap(), &[42]);
}
#[test]
fn sample_with_replaces_duplicate_field() {
let sample = Sample::new()
.with("x", Tensor::f32(&[1.0], vec![1]))
.with("x", Tensor::f32(&[2.0], vec![1]));
assert_eq!(sample.len(), 1);
assert_eq!(sample.get("x").unwrap().try_as_f32().unwrap(), &[2.0]);
}
#[test]
fn sample_insert_replaces_duplicate_field() {
let mut sample = Sample::new().with("a", Tensor::i64(&[1], vec![1]));
sample.insert("a", Tensor::i64(&[99], vec![1]));
assert_eq!(sample.len(), 1);
assert_eq!(sample.get("a").unwrap().try_as_i64().unwrap(), &[99]);
}
#[test]
fn sample_preserves_insertion_order() {
let sample = Sample::new()
.with("c", Tensor::i64(&[3], vec![1]))
.with("a", Tensor::i64(&[1], vec![1]))
.with("b", Tensor::i64(&[2], vec![1]));
let names: Vec<&str> = sample.field_names().collect();
assert_eq!(names, &["c", "a", "b"]);
}
#[test]
fn sample_remove_works() {
let mut sample = Sample::new()
.with("x", Tensor::f32(&[1.0], vec![1]))
.with("y", Tensor::i64(&[0], vec![1]));
let removed = sample.remove("x");
assert!(removed.is_some());
assert_eq!(sample.len(), 1);
assert!(sample.get("x").is_none());
assert!(sample.get("y").is_some());
}
#[test]
fn sample_remove_nonexistent_returns_none() {
let mut sample = Sample::new().with("x", Tensor::f32(&[1.0], vec![1]));
assert!(sample.remove("y").is_none());
assert_eq!(sample.len(), 1);
}
#[test]
fn sample_contains() {
let sample = Sample::new().with("x", Tensor::f32(&[1.0], vec![1]));
assert!(sample.contains("x"));
assert!(!sample.contains("y"));
}
#[test]
fn tensor_zero_copy_clone() {
let t1 = Tensor::f32(&[1.0, 2.0], vec![2]);
let t2 = t1.clone();
assert!(Arc::ptr_eq(&t1.data, &t2.data));
}
#[cfg(feature = "uring")]
#[test]
fn bytes_tensor_accepts_completed_wireshift_buffer_without_copying_to_vec() {
let pool = wireshift::BufferPool::new(8, 1).unwrap();
let mut owned = pool.acquire().unwrap();
owned.as_mut_slice()[..4].copy_from_slice(b"ring");
let completed = owned
.set_filled_len(4)
.unwrap()
.into_submitted()
.into_completed(4)
.unwrap();
let tensor = Tensor::bytes_from_completed(completed);
assert_eq!(tensor.as_bytes(), b"ring");
assert_eq!(tensor.shape(), &[4]);
}
#[test]
fn i32_tensor_round_trip() {
let tensor = Tensor::i32(&[1, -2, 3], vec![3]);
assert_eq!(tensor.dtype(), DType::I32);
assert_eq!(tensor.try_as_i32().unwrap(), &[1, -2, 3]);
}
#[test]
fn f64_tensor_round_trip() {
let tensor = Tensor::f64(&[1.5, -2.5, 3.5], vec![3]);
assert_eq!(tensor.dtype(), DType::F64);
assert_eq!(tensor.try_as_f64().unwrap(), &[1.5, -2.5, 3.5]);
}
#[test]
fn metadata() {
let sample = Sample::new()
.with("data", Tensor::u8(vec![0], vec![1]))
.with_metadata("train/001.jpg", 0);
assert_eq!(sample.metadata().unwrap().source, "train/001.jpg");
}
#[test]
fn tensor_from_bytes_validates_shape() {
let error = Tensor::from_bytes(vec![1, 2, 3], DType::I64, vec![1]).unwrap_err();
assert!(matches!(error, Error::InvalidConfig { .. }));
}
#[test]
fn misaligned_f32_view_copies_to_aligned_cache() {
let mut raw = vec![0_u8];
raw.extend_from_slice(&1.5_f32.to_le_bytes());
raw.extend_from_slice(&2.5_f32.to_le_bytes());
let cache = OnceLock::new();
let values =
ops::cast_numeric_slice::<f32, 4>(&raw[1..], &cache, f32::from_le_bytes).unwrap();
assert_eq!(values, &[1.5, 2.5]);
}
#[test]
fn misaligned_i64_view_copies_to_aligned_cache() {
let mut raw = vec![0_u8];
raw.extend_from_slice(&7_i64.to_le_bytes());
raw.extend_from_slice(&9_i64.to_le_bytes());
let cache = OnceLock::new();
let values =
ops::cast_numeric_slice::<i64, 8>(&raw[1..], &cache, i64::from_le_bytes).unwrap();
assert_eq!(values, &[7, 9]);
}
#[test]
fn bytes_tensor_shape_matches_input_length() {
let tensor = Tensor::bytes(vec![1, 2, 3, 4]);
assert_eq!(tensor.shape(), &[4]);
assert_eq!(tensor.byte_len(), 4);
}
#[test]
fn empty_sample_reports_empty() {
let sample = Sample::new();
assert!(sample.is_empty());
assert_eq!(sample.len(), 0);
}
#[test]
fn try_as_f32_wrong_dtype_returns_error() {
let tensor = Tensor::i64(&[1], vec![1]);
assert!(tensor.try_as_f32().is_err());
}
#[test]
fn try_as_f64_wrong_dtype_returns_error() {
let tensor = Tensor::f32(&[1.0], vec![1]);
assert!(tensor.try_as_f64().is_err());
}
#[test]
fn try_as_i32_wrong_dtype_returns_error() {
let tensor = Tensor::f32(&[1.0], vec![1]);
assert!(tensor.try_as_i32().is_err());
}
#[test]
fn try_as_i64_wrong_dtype_returns_error() {
let tensor = Tensor::f32(&[1.0], vec![1]);
assert!(tensor.try_as_i64().is_err());
}
#[test]
fn try_as_f32_correct_dtype_succeeds() {
let tensor = Tensor::f32(&[1.0, 2.0], vec![2]);
assert_eq!(tensor.try_as_f32().unwrap(), &[1.0, 2.0]);
}
#[test]
fn try_as_f64_correct_dtype_succeeds() {
let tensor = Tensor::f64(&[3.5, 2.75], vec![2]);
assert_eq!(tensor.try_as_f64().unwrap(), &[3.5, 2.75]);
}
#[test]
fn cast_numeric_slice_misaligned_length_returns_error() {
let data = vec![1_u8, 2, 3];
let cache = OnceLock::new();
assert!(ops::cast_numeric_slice::<f32, 4>(&data, &cache, f32::from_le_bytes).is_err());
}
#[test]
fn dtype_display_is_human_readable() {
assert_eq!(DType::F32.to_string(), "f32");
assert_eq!(DType::F64.to_string(), "f64");
assert_eq!(DType::I64.to_string(), "i64");
assert_eq!(DType::Bytes.to_string(), "bytes");
}
#[test]
fn num_elements() {
let tensor = Tensor::f32(&vec![0.0; 3 * 28 * 28], vec![3, 28, 28]);
assert_eq!(tensor.num_elements(), Some(2352));
}
}