leto 0.40.0

N-dimensional strided arrays, zero-copy views, and storage backends
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

mod broadcast;
mod contiguity;
pub(crate) mod kernels;
mod shape;
mod slice_with;
mod strides;

use serde::ser::SerializeStruct;
use serde::{Deserialize, Deserializer, Serialize, Serializer};

/// Represents an N-dimensional strided layout.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Layout<const N: usize> {
    /// The shape of the array (size of each dimension).
    pub shape: [usize; N],
    /// The stride of each dimension in elements (not bytes).
    pub strides: [isize; N],
    /// The starting offset in the storage buffer.
    pub offset: usize,
}

impl<const N: usize> Layout<N> {
    /// Create a new layout with explicit shape, strides, and offset.
    pub const fn new(shape: [usize; N], strides: [isize; N], offset: usize) -> Self {
        Self {
            shape,
            strides,
            offset,
        }
    }
}

impl<const N: usize> Serialize for Layout<N> {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        let mut state = serializer.serialize_struct("Layout", 3)?;
        state.serialize_field("shape", self.shape.as_slice())?;
        state.serialize_field("strides", self.strides.as_slice())?;
        state.serialize_field("offset", &self.offset)?;
        state.end()
    }
}

impl<'de, const N: usize> Deserialize<'de> for Layout<N> {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: Deserializer<'de>,
    {
        #[derive(Deserialize)]
        struct LayoutFields {
            shape: Vec<usize>,
            strides: Vec<isize>,
            offset: usize,
        }

        let fields = LayoutFields::deserialize(deserializer)?;
        let shape = fields.shape.try_into().map_err(|values: Vec<usize>| {
            serde::de::Error::custom(format!(
                "layout shape rank {} does not match expected rank {}",
                values.len(),
                N
            ))
        })?;
        let strides = fields.strides.try_into().map_err(|values: Vec<isize>| {
            serde::de::Error::custom(format!(
                "layout stride rank {} does not match expected rank {}",
                values.len(),
                N
            ))
        })?;

        Ok(Self {
            shape,
            strides,
            offset: fields.offset,
        })
    }
}