Skip to main content

hurray_core/layout/
coo.rs

1//! COO (Coordinate) sparse layout descriptor.
2//!
3//! Tag `0x06`. Buffer count = 2 (values + indices).
4//! See `docs/spec/layouts/coo.md`.
5
6/// Descriptor for the COO (Coordinate list) sparse layout.
7///
8/// A COO tensor stores non-zero elements as `(index_tuple, value)` pairs.
9/// It requires two buffers:
10///
11/// | Buffer | Contents |
12/// |--------|----------|
13/// | 0 — `values` | Non-zero element values; `nnz` elements of the tensor element type. |
14/// | 1 — `indices` | Index tuples as `uint64`; `nnz × rank` elements, stored row-major: `indices[i * rank + d]` is coordinate `d` of non-zero `i`. |
15///
16/// The `byte_offset` field in the common descriptor header MUST be `0` for
17/// COO tensors (no meaningful "first element at a fixed offset").
18///
19/// # Examples
20///
21/// ```
22/// use hurray_core::layout::{CooLayout, LayoutDescriptor};
23///
24/// let layout = LayoutDescriptor::Coo(CooLayout::new(3, true));
25/// assert_eq!(layout.tag(), 0x06);
26/// assert_eq!(layout.buffer_count().map(|n| n.get()), Some(2));
27/// ```
28#[derive(Debug, Clone, PartialEq, Eq, Hash)]
29#[non_exhaustive]
30pub struct CooLayout {
31    /// Number of stored (non-zero) elements. MAY be 0 for an empty sparse tensor.
32    pub nnz: u64,
33
34    /// `true` if non-zeros are stored in lexicographic index order
35    /// (dimension 0 major); `false` if no ordering guarantee is made.
36    pub is_sorted: bool,
37}
38
39impl CooLayout {
40    /// Creates a new [`CooLayout`].
41    ///
42    /// # Examples
43    ///
44    /// ```
45    /// use hurray_core::layout::CooLayout;
46    ///
47    /// let c = CooLayout::new(42, true);
48    /// assert_eq!(c.nnz, 42);
49    /// assert!(c.is_sorted);
50    /// ```
51    pub fn new(nnz: u64, is_sorted: bool) -> Self {
52        Self { nnz, is_sorted }
53    }
54}
55
56#[cfg(test)]
57mod tests {
58    use super::*;
59    use crate::layout::LayoutDescriptor;
60
61    #[test]
62    fn coo_tag_is_0x06() {
63        let layout = LayoutDescriptor::Coo(CooLayout::new(0, false));
64        assert_eq!(layout.tag(), 0x06);
65    }
66
67    #[test]
68    fn coo_buffer_count_is_2() {
69        use std::num::NonZeroU8;
70        let layout = LayoutDescriptor::Coo(CooLayout::new(100, true));
71        assert_eq!(layout.buffer_count(), Some(NonZeroU8::new(2).unwrap()));
72    }
73}