hurray_core/
shape.rs

1//! Tensor shape — rank and dimension sizes.
2//!
3//! A [`Shape`] is an ordered sequence of dimension sizes. Each size is a
4//! `u64`; the sentinel value [`DYNAMIC`] (`u64::MAX`) marks a dimension whose
5//! concrete size is not known at descriptor-write time.
6//!
7//! See `docs/spec/data-model.md` for the normative definition.
8
9use std::fmt;
10
11use crate::Error;
12
13/// Sentinel value for a **dynamic dimension** whose concrete size is not known
14/// at descriptor-write time.
15///
16/// In the wire encoding this is `0xFFFF_FFFF_FFFF_FFFF` (`UINT64_MAX`).
17/// A reader MUST NOT compute buffer sizes, strides, or element counts for a
18/// tensor containing a dynamic dimension without first resolving it to a
19/// concrete value.
20///
21/// # Examples
22///
23/// ```
24/// use hurray_core::{Shape, DYNAMIC};
25///
26/// let shape = Shape::new(vec![1, DYNAMIC, 768]).expect("valid shape");
27/// assert!(shape.has_dynamic());
28/// assert!(shape.element_count().is_none());
29/// ```
30pub const DYNAMIC: u64 = u64::MAX;
31
32/// Maximum permitted tensor rank.
33///
34/// A writer MUST NOT emit a descriptor with `rank > MAX_RANK`. A reader MUST
35/// reject any descriptor whose rank field exceeds this value.
36///
37/// The cap matches PyTorch's `MAX_DIMS = 64` and bounds the shape-array size
38/// at 512 bytes (64 × 8 bytes per dimension), enabling stack allocation on the
39/// descriptor-parsing hot path.
40///
41/// # Examples
42///
43/// ```
44/// use hurray_core::MAX_RANK;
45///
46/// assert_eq!(MAX_RANK, 64);
47/// ```
48pub const MAX_RANK: usize = 64;
49
50/// The shape of a tensor: an ordered sequence of dimension sizes.
51///
52/// Each dimension size is a [`u64`]. The special value [`DYNAMIC`] marks a
53/// dimension whose concrete size is unknown at descriptor-write time. A size
54/// of `0` denotes an empty dimension (the tensor contains no elements along
55/// that axis).
56///
57/// A [`Shape`] with no dimensions (`rank == 0`) represents a **scalar tensor**
58/// containing exactly one element.
59///
60/// # Examples
61///
62/// ```
63/// use hurray_core::{Shape, DYNAMIC};
64///
65/// // 3-D tensor with all static dimensions.
66/// let s = Shape::new(vec![3, 4, 5]).expect("valid");
67/// assert_eq!(s.rank(), 3);
68/// assert_eq!(s.element_count(), Some(60));
69///
70/// // Scalar tensor.
71/// let scalar = Shape::scalar();
72/// assert_eq!(scalar.rank(), 0);
73/// assert_eq!(scalar.element_count(), Some(1));
74///
75/// // Shape with a dynamic dimension.
76/// let dyn_shape = Shape::new(vec![1, DYNAMIC, 768]).expect("valid");
77/// assert!(dyn_shape.has_dynamic());
78/// assert!(dyn_shape.element_count().is_none());
79/// ```
80#[derive(Debug, Clone, PartialEq, Eq, Hash)]
81pub struct Shape(Vec<u64>);
82
83impl Shape {
84    /// Creates a new [`Shape`] from the given dimension sizes.
85    ///
86    /// # Errors
87    ///
88    /// Returns [`Error::RankExceedsMaximum`] if `dims.len() > MAX_RANK` (64).
89    ///
90    /// # Examples
91    ///
92    /// ```
93    /// use hurray_core::{Shape, Error};
94    ///
95    /// let s = Shape::new(vec![3, 4, 5]).expect("valid");
96    /// assert_eq!(s.rank(), 3);
97    ///
98    /// // Rank 65 must be rejected.
99    /// let too_many = Shape::new(vec![1u64; 65]);
100    /// assert!(matches!(too_many, Err(Error::RankExceedsMaximum { rank: 65, max: 64 })));
101    /// ```
102    pub fn new(dims: impl Into<Vec<u64>>) -> crate::Result<Self> {
103        let dims = dims.into();
104        if dims.len() > MAX_RANK {
105            return Err(Error::RankExceedsMaximum {
106                rank: dims.len() as u32,
107                max: MAX_RANK as u32,
108            });
109        }
110        Ok(Self(dims))
111    }
112
113    /// Returns a [`Shape`] representing a **scalar tensor** (rank 0, one element).
114    ///
115    /// # Examples
116    ///
117    /// ```
118    /// use hurray_core::Shape;
119    ///
120    /// let s = Shape::scalar();
121    /// assert_eq!(s.rank(), 0);
122    /// assert_eq!(s.dims(), &[]);
123    /// assert_eq!(s.element_count(), Some(1));
124    /// ```
125    #[inline]
126    pub fn scalar() -> Self {
127        Self(Vec::new())
128    }
129
130    /// Returns the number of dimensions (the rank) of this shape.
131    ///
132    /// # Examples
133    ///
134    /// ```
135    /// use hurray_core::Shape;
136    ///
137    /// assert_eq!(Shape::scalar().rank(), 0);
138    /// assert_eq!(Shape::new(vec![3, 4]).unwrap().rank(), 2);
139    /// ```
140    #[inline]
141    pub fn rank(&self) -> usize {
142        self.0.len()
143    }
144
145    /// Returns the dimension sizes as a slice.
146    ///
147    /// # Examples
148    ///
149    /// ```
150    /// use hurray_core::Shape;
151    ///
152    /// let s = Shape::new(vec![2, 3]).unwrap();
153    /// assert_eq!(s.dims(), &[2, 3]);
154    /// ```
155    #[inline]
156    pub fn dims(&self) -> &[u64] {
157        &self.0
158    }
159
160    /// Returns the total number of logical elements, or `None` if any
161    /// dimension is [`DYNAMIC`] or if the product overflows `u64`.
162    ///
163    /// The product of an empty sequence (scalar) is `1`. A tensor with any
164    /// zero-size dimension has `element_count == Some(0)`.
165    ///
166    /// A reader MUST NOT use this value to compute buffer sizes or strides
167    /// until all dynamic dimensions have been resolved.
168    ///
169    /// # Examples
170    ///
171    /// ```
172    /// use hurray_core::{Shape, DYNAMIC};
173    ///
174    /// // Static shape.
175    /// assert_eq!(Shape::new(vec![3, 4, 5]).unwrap().element_count(), Some(60));
176    ///
177    /// // Scalar: product of empty sequence is 1.
178    /// assert_eq!(Shape::scalar().element_count(), Some(1));
179    ///
180    /// // Empty tensor: any zero dimension makes the count 0.
181    /// assert_eq!(Shape::new(vec![3, 0, 5]).unwrap().element_count(), Some(0));
182    ///
183    /// // Dynamic dimension: result is unknown.
184    /// assert_eq!(Shape::new(vec![1, DYNAMIC, 768]).unwrap().element_count(), None);
185    /// ```
186    pub fn element_count(&self) -> Option<u64> {
187        let mut product: u64 = 1;
188        for &dim in &self.0 {
189            if dim == DYNAMIC {
190                return None;
191            }
192            product = product.checked_mul(dim)?;
193        }
194        Some(product)
195    }
196
197    /// Returns `true` if any dimension has size `0`.
198    ///
199    /// An empty tensor is valid; its data buffer has size 0 bytes. Note that
200    /// a [`DYNAMIC`] dimension does not make a tensor empty — its size is
201    /// unknown, not zero.
202    ///
203    /// # Examples
204    ///
205    /// ```
206    /// use hurray_core::{Shape, DYNAMIC};
207    ///
208    /// assert!(Shape::new(vec![3, 0, 5]).unwrap().is_empty_tensor());
209    /// assert!(!Shape::new(vec![3, 4, 5]).unwrap().is_empty_tensor());
210    /// // Dynamic is not empty.
211    /// assert!(!Shape::new(vec![1, DYNAMIC, 768]).unwrap().is_empty_tensor());
212    /// ```
213    #[inline]
214    pub fn is_empty_tensor(&self) -> bool {
215        self.0.contains(&0)
216    }
217
218    /// Returns `true` if any dimension equals [`DYNAMIC`].
219    ///
220    /// # Examples
221    ///
222    /// ```
223    /// use hurray_core::{Shape, DYNAMIC};
224    ///
225    /// assert!(Shape::new(vec![1, DYNAMIC, 768]).unwrap().has_dynamic());
226    /// assert!(!Shape::new(vec![1, 128, 768]).unwrap().has_dynamic());
227    /// ```
228    #[inline]
229    pub fn has_dynamic(&self) -> bool {
230        self.0.contains(&DYNAMIC)
231    }
232}
233
234impl fmt::Display for Shape {
235    /// Formats the shape as a bracket-enclosed, comma-separated list of sizes.
236    ///
237    /// Dynamic dimensions are shown as `?`. An empty (scalar) shape is `[]`.
238    ///
239    /// # Examples
240    ///
241    /// ```
242    /// use hurray_core::{Shape, DYNAMIC};
243    ///
244    /// assert_eq!(Shape::scalar().to_string(), "[]");
245    /// assert_eq!(Shape::new(vec![3, 4, 5]).unwrap().to_string(), "[3, 4, 5]");
246    /// assert_eq!(Shape::new(vec![1, DYNAMIC, 768]).unwrap().to_string(), "[1, ?, 768]");
247    /// ```
248    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
249        f.write_str("[")?;
250        for (i, &dim) in self.0.iter().enumerate() {
251            if i > 0 {
252                f.write_str(", ")?;
253            }
254            if dim == DYNAMIC {
255                f.write_str("?")?;
256            } else {
257                write!(f, "{dim}")?;
258            }
259        }
260        f.write_str("]")
261    }
262}
263
264#[cfg(test)]
265mod tests {
266    use super::*;
267
268    // ── Shape::new ───────────────────────────────────────────────────────────
269
270    #[test]
271    fn new_rank_0_succeeds() {
272        let s = Shape::new(vec![]).expect("rank 0 is valid");
273        assert_eq!(s.rank(), 0);
274    }
275
276    #[test]
277    fn new_rank_1_succeeds() {
278        let s = Shape::new(vec![42]).expect("rank 1 is valid");
279        assert_eq!(s.rank(), 1);
280        assert_eq!(s.dims(), &[42]);
281    }
282
283    #[test]
284    fn new_rank_64_succeeds() {
285        let s = Shape::new(vec![1u64; 64]).expect("rank 64 is valid (MAX_RANK)");
286        assert_eq!(s.rank(), 64);
287    }
288
289    /// Spec § data-model: a writer MUST NOT emit a descriptor with rank > MAX_RANK;
290    /// a reader MUST reject any descriptor whose rank field exceeds this value.
291    #[test]
292    fn new_rank_65_returns_rank_exceeds_maximum() {
293        let result = Shape::new(vec![1u64; 65]);
294        assert!(
295            matches!(result, Err(Error::RankExceedsMaximum { rank: 65, max: 64 })),
296            "expected RankExceedsMaximum for rank 65, got {result:?}"
297        );
298    }
299
300    #[test]
301    fn new_rank_100_returns_rank_exceeds_maximum() {
302        let result = Shape::new(vec![1u64; 100]);
303        assert!(matches!(
304            result,
305            Err(Error::RankExceedsMaximum { rank: 100, max: 64 })
306        ));
307    }
308
309    // ── Shape::scalar ────────────────────────────────────────────────────────
310
311    #[test]
312    fn scalar_has_rank_0() {
313        assert_eq!(Shape::scalar().rank(), 0);
314    }
315
316    #[test]
317    fn scalar_dims_is_empty() {
318        assert_eq!(Shape::scalar().dims(), &[] as &[u64]);
319    }
320
321    #[test]
322    fn scalar_element_count_is_one() {
323        assert_eq!(Shape::scalar().element_count(), Some(1));
324    }
325
326    #[test]
327    fn scalar_is_not_empty_tensor() {
328        assert!(!Shape::scalar().is_empty_tensor());
329    }
330
331    #[test]
332    fn scalar_has_no_dynamic() {
333        assert!(!Shape::scalar().has_dynamic());
334    }
335
336    // ── rank ─────────────────────────────────────────────────────────────────
337
338    #[test]
339    fn rank_matches_dims_len() {
340        for n in [0, 1, 2, 10, 64] {
341            let s = Shape::new(vec![3u64; n]).expect("valid");
342            assert_eq!(s.rank(), n, "rank mismatch for n={n}");
343        }
344    }
345
346    // ── dims ─────────────────────────────────────────────────────────────────
347
348    #[test]
349    fn dims_returns_original_slice() {
350        let input = vec![3u64, 4, 5];
351        let s = Shape::new(input.clone()).expect("valid");
352        assert_eq!(s.dims(), input.as_slice());
353    }
354
355    #[test]
356    fn dims_is_empty_for_scalar() {
357        assert!(Shape::scalar().dims().is_empty());
358    }
359
360    // ── element_count ────────────────────────────────────────────────────────
361
362    #[test]
363    fn element_count_static_shape_3x4x5() {
364        let s = Shape::new(vec![3, 4, 5]).expect("valid");
365        assert_eq!(s.element_count(), Some(60));
366    }
367
368    #[test]
369    fn element_count_1d_shape() {
370        assert_eq!(Shape::new(vec![7]).unwrap().element_count(), Some(7));
371    }
372
373    #[test]
374    fn element_count_scalar_is_one() {
375        assert_eq!(Shape::scalar().element_count(), Some(1));
376    }
377
378    #[test]
379    fn element_count_with_zero_dim_is_zero() {
380        assert_eq!(Shape::new(vec![3, 0, 5]).unwrap().element_count(), Some(0));
381    }
382
383    #[test]
384    fn element_count_all_zero_dims_is_zero() {
385        assert_eq!(Shape::new(vec![0, 0]).unwrap().element_count(), Some(0));
386    }
387
388    #[test]
389    fn element_count_with_dynamic_is_none() {
390        assert_eq!(
391            Shape::new(vec![1, DYNAMIC, 768]).unwrap().element_count(),
392            None
393        );
394    }
395
396    #[test]
397    fn element_count_all_dynamic_is_none() {
398        assert_eq!(
399            Shape::new(vec![DYNAMIC, DYNAMIC]).unwrap().element_count(),
400            None
401        );
402    }
403
404    /// Overflow during product computation must return None, not panic.
405    #[test]
406    fn element_count_overflow_returns_none() {
407        // u64::MAX / 2 + 1 times two dims = overflow
408        let big = u64::MAX / 2 + 1;
409        let s = Shape::new(vec![big, 2]).expect("rank 2 is valid");
410        assert_eq!(s.element_count(), None);
411    }
412
413    // ── is_empty_tensor ──────────────────────────────────────────────────────
414
415    #[test]
416    fn is_empty_tensor_true_when_one_dim_is_zero() {
417        assert!(Shape::new(vec![3, 0, 5]).unwrap().is_empty_tensor());
418    }
419
420    #[test]
421    fn is_empty_tensor_true_when_all_dims_zero() {
422        assert!(Shape::new(vec![0]).unwrap().is_empty_tensor());
423    }
424
425    #[test]
426    fn is_empty_tensor_false_for_static_nonzero_shape() {
427        assert!(!Shape::new(vec![3, 4, 5]).unwrap().is_empty_tensor());
428    }
429
430    /// A DYNAMIC dimension is unknown, not zero — the tensor is not empty.
431    #[test]
432    fn is_empty_tensor_false_for_dynamic_dim() {
433        assert!(!Shape::new(vec![1, DYNAMIC, 768]).unwrap().is_empty_tensor());
434    }
435
436    #[test]
437    fn is_empty_tensor_false_for_scalar() {
438        assert!(!Shape::scalar().is_empty_tensor());
439    }
440
441    // ── has_dynamic ──────────────────────────────────────────────────────────
442
443    #[test]
444    fn has_dynamic_true_when_one_dim_is_dynamic() {
445        assert!(Shape::new(vec![1, DYNAMIC, 768]).unwrap().has_dynamic());
446    }
447
448    #[test]
449    fn has_dynamic_true_when_all_dims_dynamic() {
450        assert!(Shape::new(vec![DYNAMIC, DYNAMIC]).unwrap().has_dynamic());
451    }
452
453    #[test]
454    fn has_dynamic_false_for_fully_static_shape() {
455        assert!(!Shape::new(vec![1, 128, 768]).unwrap().has_dynamic());
456    }
457
458    #[test]
459    fn has_dynamic_false_for_scalar() {
460        assert!(!Shape::scalar().has_dynamic());
461    }
462
463    // ── Display ──────────────────────────────────────────────────────────────
464
465    #[test]
466    fn display_scalar_is_empty_brackets() {
467        assert_eq!(Shape::scalar().to_string(), "[]");
468    }
469
470    #[test]
471    fn display_1d_shape() {
472        assert_eq!(Shape::new(vec![5]).unwrap().to_string(), "[5]");
473    }
474
475    #[test]
476    fn display_3d_static_shape() {
477        assert_eq!(Shape::new(vec![3, 4, 5]).unwrap().to_string(), "[3, 4, 5]");
478    }
479
480    #[test]
481    fn display_dynamic_dimension_shown_as_question_mark() {
482        assert_eq!(
483            Shape::new(vec![1, DYNAMIC, 768]).unwrap().to_string(),
484            "[1, ?, 768]"
485        );
486    }
487
488    #[test]
489    fn display_all_dynamic() {
490        assert_eq!(
491            Shape::new(vec![DYNAMIC, DYNAMIC]).unwrap().to_string(),
492            "[?, ?]"
493        );
494    }
495
496    #[test]
497    fn display_zero_dimension() {
498        assert_eq!(Shape::new(vec![3, 0, 5]).unwrap().to_string(), "[3, 0, 5]");
499    }
500
501    // ── DYNAMIC and MAX_RANK constants ───────────────────────────────────────
502
503    #[test]
504    fn dynamic_constant_is_u64_max() {
505        assert_eq!(DYNAMIC, u64::MAX);
506    }
507
508    #[test]
509    fn max_rank_constant_is_64() {
510        assert_eq!(MAX_RANK, 64);
511    }
512}
hurray_core/shape.rs

hurray_core/
shape.rs
