1use super::Shape;
2use crate::Result;
3use crate::{DTypeKind, Device, Error};
4use std::fmt::Display;
5
6pub struct DimCoordinates {
7 shape: Vec<usize>,
8 current: Vec<usize>,
9 done: bool,
10}
11
12impl DimCoordinates {
13 pub fn from_shape(shape: &Shape) -> Self {
14 let rank = shape.rank();
15 Self { shape: shape.dims().to_vec(), current: vec![0; rank], done: shape.is_scalar() }
16 }
17}
18
19impl Iterator for DimCoordinates {
20 type Item = Vec<usize>;
21
22 fn next(&mut self) -> Option<Self::Item> {
23 if self.done {
24 return None;
25 }
26
27 let result = self.current.clone();
28
29 for i in (0..self.current.len()).rev() {
30 self.current[i] += 1;
31 if self.current[i] < self.shape[i] {
32 break;
33 } else {
34 self.current[i] = 0;
35 if i == 0 {
36 self.done = true;
37 }
38 }
39 }
40
41 Some(result)
42 }
43}
44
45pub struct DimNCoordinates<const N: usize> {
46 shape: [usize; N],
47 current: [usize; N],
48 done: bool,
49}
50
51impl<const N: usize> DimNCoordinates<N> {
52 pub fn from_shape(from_shape: &Shape) -> Result<Self> {
53 if from_shape.rank() == N {
54 let mut shape = [0usize; N];
55 for i in 0..N {
56 shape[i] = from_shape.dims()[i];
57 }
58
59 let current = [0usize; N];
60 Ok(Self { shape, current, done: N == 0 })
61 } else {
62 Err(Error::UnexpectedNumberOfDims { expected: N, got: from_shape.rank(), shape: Shape::from(from_shape.dims()) })?
63 }
64 }
65}
66
67impl<const N: usize> Iterator for DimNCoordinates<N> {
68 type Item = [usize; N];
69 fn next(&mut self) -> Option<Self::Item> {
70 if self.done {
71 return None;
72 }
73
74 let result = self.current;
75
76 for i in (0..N).rev() {
77 self.current[i] += 1;
78 if self.current[i] < self.shape[i] {
79 break;
80 } else {
81 self.current[i] = 0;
82 if i == 0 {
83 self.done = true;
84 }
85 }
86 }
87
88 Some(result)
89 }
90}
91
92impl<const C: usize> From<&[usize; C]> for Shape {
93 fn from(dims: &[usize; C]) -> Self {
94 Self(dims.to_vec())
95 }
96}
97
98impl From<Vec<usize>> for Shape {
99 fn from(dims: Vec<usize>) -> Self {
100 Self(dims)
101 }
102}
103
104impl From<&Vec<usize>> for Shape {
105 fn from(dims: &Vec<usize>) -> Self {
106 Self(dims.clone())
107 }
108}
109
110impl From<&[usize]> for Shape {
111 fn from(dims: &[usize]) -> Self {
112 Self(dims.to_vec())
113 }
114}
115
116impl From<&Shape> for Shape {
117 fn from(shape: &Shape) -> Self {
118 Self(shape.0.to_vec())
119 }
120}
121
122impl From<usize> for Shape {
123 fn from(d1: usize) -> Self {
124 Self([d1].to_vec())
125 }
126}
127
128impl From<()> for Shape {
129 fn from(_: ()) -> Self {
130 Self(vec![])
131 }
132}
133
134impl std::fmt::Display for Shape {
135 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
136 write!(f, "(")?;
137 for (i, dim) in self.0.iter().enumerate() {
138 if i > 0 {
139 write!(f, ", ")?;
140 }
141 write!(f, "{}", dim)?;
142 }
143 if self.0.len() == 1 {
144 write!(f, ",")?;
145 }
146 write!(f, ")")
147 }
148}
149
150macro_rules! impl_from_tuple {
151 ($tuple:ty, $($index:tt),+) => {
152 impl From<$tuple> for Shape {
153 fn from(d: $tuple) -> Self {
154 Self([$(d.$index,)+].to_vec())
155 }
156 }
157 };
158}
159
160impl_from_tuple!((usize,), 0);
161impl_from_tuple!((usize, usize), 0, 1);
162impl_from_tuple!((usize, usize, usize), 0, 1, 2);
163impl_from_tuple!((usize, usize, usize, usize), 0, 1, 2, 3);
164impl_from_tuple!((usize, usize, usize, usize, usize), 0, 1, 2, 3, 4);
165impl_from_tuple!((usize, usize, usize, usize, usize, usize), 0, 1, 2, 3, 4, 5);
166
167#[derive(Debug, Clone, Copy, PartialEq, Eq)]
168pub enum D {
169 Minus1,
170 Minus2,
171 Minus(usize),
172 Index(usize),
173}
174
175impl D {
176 pub fn to_real_index(&self, size: usize, op: &'static str) -> Result<usize> {
177 match self {
178 Self::Minus1 if size >= 1 => Ok(size - 1),
179 Self::Minus2 if size >= 2 => Ok(size - 2),
180 Self::Minus(u) if *u > 0 && size >= *u => Ok(size - *u),
181 Self::Index(u) if *u < size => Ok(*u),
182 _ => Err(crate::Error::DimSizeOutOfRange { size, op })?,
183 }
184 }
185}
186
187impl Display for D {
188 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
189 match self {
190 Self::Minus(n) => writeln!(f, "-{}", n),
191 Self::Minus1 => writeln!(f, "-1"),
192 Self::Minus2 => writeln!(f, "-2"),
193 Self::Index(n) => writeln!(f, "{}", n),
194 }
195 }
196}
197
198impl D {
199 fn out_of_range(&self, shape: &Shape, op: &'static str) -> Error {
200 let dim = match self {
201 Self::Minus1 => -1,
202 Self::Minus2 => -2,
203 Self::Minus(u) => -(*u as i32),
204 Self::Index(u) => *u as i32,
205 };
206 Error::DimOutOfRange { shape: shape.clone(), dim, op }
207 }
208}
209
210macro_rules! extract_dims {
211 ($fn_name:ident, $cnt:tt, $dims:expr, $out_type:ty) => {
212 pub fn $fn_name(dims: &[usize]) -> Result<$out_type> {
213 if dims.len() != $cnt {
214 Err(Error::UnexpectedNumberOfDims { expected: $cnt, got: dims.len(), shape: Shape::from(dims) })?
215 } else {
216 Ok($dims(dims))
217 }
218 }
219
220 impl Shape {
221 pub fn $fn_name(&self) -> Result<$out_type> {
222 $fn_name(self.0.as_slice())
223 }
224 }
225
226 impl<D: Device, K: DTypeKind<D>> crate::Tensor<D, K> {
227 pub fn $fn_name(&self) -> Result<$out_type> {
228 self.shape().$fn_name()
229 }
230 }
231
232 impl std::convert::TryInto<$out_type> for Shape {
233 type Error = crate::Error;
234 fn try_into(self) -> crate::Result<$out_type> {
235 self.$fn_name()
236 }
237 }
238 };
239}
240
241extract_dims!(dims0, 0, |_: &[usize]| (), ());
242extract_dims!(dims1, 1, |d: &[usize]| d[0], usize);
243extract_dims!(dims2, 2, |d: &[usize]| (d[0], d[1]), (usize, usize));
244extract_dims!(dims3, 3, |d: &[usize]| (d[0], d[1], d[2]), (usize, usize, usize));
245extract_dims!(dims4, 4, |d: &[usize]| (d[0], d[1], d[2], d[3]), (usize, usize, usize, usize));
246extract_dims!(dims5, 5, |d: &[usize]| (d[0], d[1], d[2], d[3], d[4]), (usize, usize, usize, usize, usize));
247
248pub trait Dim: Copy {
249 fn to_index(&self, shape: &Shape, op: &'static str) -> Result<usize>;
250 fn to_index_plus_one(&self, shape: &Shape, op: &'static str) -> Result<usize>;
251}
252
253impl Dim for usize {
254 fn to_index(&self, shape: &Shape, op: &'static str) -> Result<usize> {
255 let dim = *self;
256 if dim >= shape.rank() { Err(Error::DimOutOfRange { shape: shape.clone(), dim: dim as i32, op })? } else { Ok(dim) }
257 }
258
259 fn to_index_plus_one(&self, shape: &Shape, op: &'static str) -> Result<usize> {
260 let dim = *self;
261 if dim > shape.rank() { Err(Error::DimOutOfRange { shape: shape.clone(), dim: dim as i32, op })? } else { Ok(dim) }
262 }
263}
264
265impl Dim for D {
266 fn to_index(&self, shape: &Shape, op: &'static str) -> Result<usize> {
267 let rank = shape.rank();
268 match self {
269 Self::Minus1 if rank >= 1 => Ok(rank - 1),
270 Self::Minus2 if rank >= 2 => Ok(rank - 2),
271 Self::Minus(u) if *u > 0 && rank >= *u => Ok(rank - *u),
272 Self::Index(u) => u.to_index(shape, op),
273 _ => Err(self.out_of_range(shape, op))?,
274 }
275 }
276
277 fn to_index_plus_one(&self, shape: &Shape, op: &'static str) -> Result<usize> {
278 let rank = shape.rank();
279 match self {
280 Self::Minus1 => Ok(rank),
281 Self::Minus2 if rank >= 1 => Ok(rank - 1),
282 Self::Minus(u) if *u > 0 && rank + 1 >= *u => Ok(rank + 1 - *u),
283 Self::Index(u) => u.to_index_plus_one(shape, op),
284 _ => Err(self.out_of_range(shape, op))?,
285 }
286 }
287}
288
289pub trait Dims {
290 fn to_indexes(self, shape: &Shape, op: &'static str) -> Result<Vec<usize>>;
291
292 fn check_indexes(dims: &[usize], shape: &Shape, op: &'static str) -> Result<()> {
293 for (i, &dim) in dims.iter().enumerate() {
294 if dims[..i].contains(&dim) {
295 return Err(Error::DuplicateDimIndex { shape: shape.clone(), dims: dims.to_vec(), op })?;
296 }
297 if dim >= shape.rank() {
298 return Err(Error::DimOutOfRange { shape: shape.clone(), dim: dim as i32, op })?;
299 }
300 }
301 Ok(())
302 }
303}
304
305impl Dims for Vec<usize> {
306 fn to_indexes(self, shape: &Shape, op: &'static str) -> Result<Vec<usize>> {
307 Self::check_indexes(&self, shape, op)?;
308 Ok(self)
309 }
310}
311
312impl<const N: usize> Dims for [usize; N] {
313 fn to_indexes(self, shape: &Shape, op: &'static str) -> Result<Vec<usize>> {
314 Self::check_indexes(&self, shape, op)?;
315 Ok(self.to_vec())
316 }
317}
318
319impl Dims for &[usize] {
320 fn to_indexes(self, shape: &Shape, op: &'static str) -> Result<Vec<usize>> {
321 Self::check_indexes(self, shape, op)?;
322 Ok(self.to_vec())
323 }
324}
325
326impl Dims for () {
327 fn to_indexes(self, _: &Shape, _: &'static str) -> Result<Vec<usize>> {
328 Ok(vec![])
329 }
330}
331
332impl<Di: Dim + Sized> Dims for Di {
333 fn to_indexes(self, shape: &Shape, op: &'static str) -> Result<Vec<usize>> {
334 let dim = self.to_index(shape, op)?;
335 Ok([dim].to_vec())
336 }
337}
338
339impl<D1: Dim, D2: Dim> Dims for (D1, D2) {
340 fn to_indexes(self, shape: &Shape, op: &'static str) -> Result<Vec<usize>> {
341 let d0 = self.0.to_index(shape, op)?;
342 let d1 = self.1.to_index(shape, op)?;
343 Ok([d0, d1].to_vec())
344 }
345}
346
347impl<D1: Dim, D2: Dim, D3: Dim> Dims for (D1, D2, D3) {
348 fn to_indexes(self, shape: &Shape, op: &'static str) -> Result<Vec<usize>> {
349 let d0 = self.0.to_index(shape, op)?;
350 let d1 = self.1.to_index(shape, op)?;
351 let d2 = self.2.to_index(shape, op)?;
352 Ok([d0, d1, d2].to_vec())
353 }
354}
355
356impl<D1: Dim, D2: Dim, D3: Dim, D4: Dim> Dims for (D1, D2, D3, D4) {
357 fn to_indexes(self, shape: &Shape, op: &'static str) -> Result<Vec<usize>> {
358 let d0 = self.0.to_index(shape, op)?;
359 let d1 = self.1.to_index(shape, op)?;
360 let d2 = self.2.to_index(shape, op)?;
361 let d3 = self.3.to_index(shape, op)?;
362 Ok([d0, d1, d2, d3].to_vec())
363 }
364}
365
366impl<D1: Dim, D2: Dim, D3: Dim, D4: Dim, D5: Dim> Dims for (D1, D2, D3, D4, D5) {
367 fn to_indexes(self, shape: &Shape, op: &'static str) -> Result<Vec<usize>> {
368 let d0 = self.0.to_index(shape, op)?;
369 let d1 = self.1.to_index(shape, op)?;
370 let d2 = self.2.to_index(shape, op)?;
371 let d3 = self.3.to_index(shape, op)?;
372 let d4 = self.4.to_index(shape, op)?;
373 Ok([d0, d1, d2, d3, d4].to_vec())
374 }
375}