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
use super::TensorData;
use crate::tensor::{BoolStore, DType, QuantLevel, QuantMode, QuantScheme, QuantValue, QuantizedBytes};
use half::{bf16, f16};
use crate::tensor::element::{Element, ElementConversion};
use alloc::boxed::Box;
use alloc::vec::Vec;
impl TensorData {
/// Returns an iterator over the values of the tensor data.
pub fn iter<E: Element>(&self) -> Box<dyn Iterator<Item = E> + '_> {
if E::dtype() == self.dtype {
Box::new(bytemuck::checked::cast_slice(&self.bytes).iter().copied())
} else {
match self.dtype {
DType::I8 => Box::new(
bytemuck::checked::cast_slice(&self.bytes)
.iter()
.map(|e: &i8| e.elem::<E>()),
),
DType::I16 => Box::new(
bytemuck::checked::cast_slice(&self.bytes)
.iter()
.map(|e: &i16| e.elem::<E>()),
),
DType::I32 => Box::new(
bytemuck::checked::cast_slice(&self.bytes)
.iter()
.map(|e: &i32| e.elem::<E>()),
),
DType::I64 => Box::new(
bytemuck::checked::cast_slice(&self.bytes)
.iter()
.map(|e: &i64| e.elem::<E>()),
),
DType::U8 => Box::new(self.bytes.iter().map(|e| e.elem::<E>())),
DType::U16 => Box::new(
bytemuck::checked::cast_slice(&self.bytes)
.iter()
.map(|e: &u16| e.elem::<E>()),
),
DType::U32 => Box::new(
bytemuck::checked::cast_slice(&self.bytes)
.iter()
.map(|e: &u32| e.elem::<E>()),
),
DType::U64 => Box::new(
bytemuck::checked::cast_slice(&self.bytes)
.iter()
.map(|e: &u64| e.elem::<E>()),
),
DType::BF16 => Box::new(
bytemuck::checked::cast_slice(&self.bytes)
.iter()
.map(|e: &bf16| e.elem::<E>()),
),
DType::F16 => Box::new(
bytemuck::checked::cast_slice(&self.bytes)
.iter()
.map(|e: &f16| e.elem::<E>()),
),
DType::F32 | DType::Flex32 => Box::new(
bytemuck::checked::cast_slice(&self.bytes)
.iter()
.map(|e: &f32| e.elem::<E>()),
),
DType::F64 => Box::new(
bytemuck::checked::cast_slice(&self.bytes)
.iter()
.map(|e: &f64| e.elem::<E>()),
),
// bool is a byte value equal to either 0 or 1
DType::Bool(BoolStore::Native) | DType::Bool(BoolStore::U8) => {
Box::new(self.bytes.iter().map(|e| e.elem::<E>()))
}
DType::Bool(BoolStore::U32) => Box::new(
bytemuck::checked::cast_slice(&self.bytes)
.iter()
.map(|e: &u32| e.elem::<E>()),
),
DType::QFloat(scheme) => match scheme {
QuantScheme {
level: QuantLevel::Tensor | QuantLevel::Block(_),
mode: QuantMode::Symmetric,
value:
QuantValue::Q8F
| QuantValue::Q8S
// Represent sub-byte values as i8
| QuantValue::Q4F
| QuantValue::Q4S
| QuantValue::Q2F
| QuantValue::Q2S,
..
} => {
// Quantized int8 values
let q_bytes = QuantizedBytes {
bytes: self.bytes.clone(),
scheme,
num_elements: self.num_elements(),
};
let (values, _) = q_bytes.into_vec_i8_with_shape(&self.shape);
Box::new(
values
.iter()
.map(|e: &i8| e.elem::<E>())
.collect::<Vec<_>>()
.into_iter(),
)
}
QuantScheme {
level: QuantLevel::Tensor | QuantLevel::Block(_),
mode: QuantMode::Symmetric,
value:
QuantValue::E4M3 | QuantValue::E5M2 | QuantValue::E2M1,
..
} => {
unimplemented!("Not yet implemented for iteration");
}
},
}
}
}
}