xn 0.1.1

Another minimalist deep-learning framework optimized for inference.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125

use std::collections::HashMap;
use xn::{CPU, CpuTensor, DType, Tensor, TypedTensor, safetensors};

#[test]
fn save_and_load_single_tensor() {
    let path = std::env::temp_dir().join("test_single.safetensors");
    let t: CpuTensor<f32> = Tensor::from_vec(vec![1.0, 2.0, 3.0, 4.0], (2, 2), &CPU).unwrap();
    t.save_safetensors("t", &path).unwrap();

    let tensors = safetensors::load_from_file(&path, &CPU).unwrap();
    assert_eq!(tensors.len(), 1);
    let loaded = match tensors.get("t").unwrap() {
        TypedTensor::F32(t) => t,
        _ => panic!("expected F32"),
    };
    assert_eq!(loaded.dims(), &[2, 2]);
    assert_eq!(loaded.to_vec().unwrap(), vec![1.0, 2.0, 3.0, 4.0]);

    std::fs::remove_file(&path).unwrap();
}

#[test]
fn save_and_load_multiple_tensors() {
    let path = std::env::temp_dir().join("test_multi.safetensors");
    let a: CpuTensor<f32> = Tensor::from_vec(vec![1.0, 2.0, 3.0], (1, 3), &CPU).unwrap();
    let b: CpuTensor<i64> = Tensor::from_vec(vec![10, 20], (2,), &CPU).unwrap();

    let mut map: HashMap<String, TypedTensor<_>> = HashMap::new();
    map.insert("a".to_string(), TypedTensor::F32(a));
    map.insert("b".to_string(), TypedTensor::I64(b));
    safetensors::save(&map, &path).unwrap();

    let loaded = safetensors::load_from_file(&path, &CPU).unwrap();
    assert_eq!(loaded.len(), 2);

    let la = match loaded.get("a").unwrap() {
        TypedTensor::F32(t) => t,
        _ => panic!("expected F32"),
    };
    assert_eq!(la.dims(), &[1, 3]);
    assert_eq!(la.to_vec().unwrap(), vec![1.0, 2.0, 3.0]);

    let lb = match loaded.get("b").unwrap() {
        TypedTensor::I64(t) => t,
        _ => panic!("expected I64"),
    };
    assert_eq!(lb.dims(), &[2]);
    assert_eq!(lb.to_vec().unwrap(), vec![10, 20]);

    std::fs::remove_file(&path).unwrap();
}

#[test]
fn save_with_str_keys() {
    let path = std::env::temp_dir().join("test_str_keys.safetensors");
    let t: CpuTensor<f32> = Tensor::from_vec(vec![5.0, 6.0], (2,), &CPU).unwrap();

    let map: HashMap<&str, TypedTensor<_>> = [("x", TypedTensor::F32(t))].into_iter().collect();
    safetensors::save(&map, &path).unwrap();

    let loaded = safetensors::load_from_file(&path, &CPU).unwrap();
    let lx = match loaded.get("x").unwrap() {
        TypedTensor::F32(t) => t,
        _ => panic!("expected F32"),
    };
    assert_eq!(lx.to_vec().unwrap(), vec![5.0, 6.0]);

    std::fs::remove_file(&path).unwrap();
}

#[test]
fn roundtrip_preserves_dtypes() {
    let path = std::env::temp_dir().join("test_dtypes.safetensors");

    let f16_t: CpuTensor<half::f16> =
        Tensor::from_vec(vec![half::f16::from_f32(1.0)], (1,), &CPU).unwrap();
    let bf16_t: CpuTensor<half::bf16> =
        Tensor::from_vec(vec![half::bf16::from_f32(2.0)], (1,), &CPU).unwrap();
    let u8_t: CpuTensor<u8> = Tensor::from_vec(vec![42], (1,), &CPU).unwrap();

    let mut map: HashMap<String, TypedTensor<_>> = HashMap::new();
    map.insert("f16".to_string(), TypedTensor::F16(f16_t));
    map.insert("bf16".to_string(), TypedTensor::BF16(bf16_t));
    map.insert("u8".to_string(), TypedTensor::U8(u8_t));
    safetensors::save(&map, &path).unwrap();

    let loaded = safetensors::load_from_file(&path, &CPU).unwrap();

    assert_eq!(loaded.get("f16").unwrap().dtype(), DType::F16);
    assert_eq!(loaded.get("bf16").unwrap().dtype(), DType::BF16);
    assert_eq!(loaded.get("u8").unwrap().dtype(), DType::U8);

    let lf16 = match loaded.get("f16").unwrap() {
        TypedTensor::F16(t) => t,
        _ => panic!("expected F16"),
    };
    assert_eq!(lf16.to_vec().unwrap(), vec![half::f16::from_f32(1.0)]);

    let lu8 = match loaded.get("u8").unwrap() {
        TypedTensor::U8(t) => t,
        _ => panic!("expected U8"),
    };
    assert_eq!(lu8.to_vec().unwrap(), vec![42]);

    std::fs::remove_file(&path).unwrap();
}

#[test]
fn save_load_buffer_roundtrip() {
    let path = std::env::temp_dir().join("test_buffer.safetensors");
    let t: CpuTensor<f32> =
        Tensor::from_vec(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0], (2, 3), &CPU).unwrap();
    t.save_safetensors("t", &path).unwrap();

    let bytes = std::fs::read(&path).unwrap();
    let loaded = safetensors::load_from_buffer(&bytes, &CPU).unwrap();
    let lt = match loaded.get("t").unwrap() {
        TypedTensor::F32(t) => t,
        _ => panic!("expected F32"),
    };
    assert_eq!(lt.dims(), &[2, 3]);
    assert_eq!(lt.to_vec().unwrap(), vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]);

    std::fs::remove_file(&path).unwrap();
}