burn_tensor/
shape_macros.rs

1//! Shape assertion macros. Parsing and expansion live in burn-derive; these wrappers forward
2//! `$crate` so the expansion can name [`Tensor::dims`](crate::Tensor::dims) by path, which
3//! makes anything other than a `Tensor` a type error.
4//!
5//! Inspired by the `burn-contracts` crate by Crutcher Dunnavant.
6
7/// Asserts a tensor's rank, at compile time unless the pattern has `..`, and its axis sizes at
8/// runtime.
9///
10/// The pattern has one slot per axis. A slot is either `_`, which skips the axis, or any `usize`
11/// expression the axis must equal: a name from `dims()`, a config field, a literal, arithmetic
12/// on any of those. To name axes, destructure `tensor.dims()` first and use the names as slots.
13///
14/// The pattern length must equal the tensor rank, otherwise the call does not compile. One
15/// slot may be `..`, which stands for any number of axes; the pattern then fixes only a
16/// minimum rank, checked at runtime, and fits a function generic over `const D: usize`. A
17/// pattern of just `[..]` checks nothing beyond the argument being a tensor. The tensor is
18/// borrowed, not moved, and every expression is evaluated once. Checks stay on in release
19/// builds; see [`debug_assert_shape!`](crate::debug_assert_shape) for the debug-only variant.
20///
21/// ```
22/// use burn_tensor::{Tensor, assert_shape};
23///
24/// let device = Default::default();
25/// let x = Tensor::<3>::zeros([2, 5, 256], &device);
26/// let mask = Tensor::<2>::zeros([2, 5], &device);
27/// let patch = Tensor::<2>::zeros([6, 2], &device);
28///
29/// let [batch_size, seq_length, _] = x.dims();
30/// let flat = x.clone().reshape([batch_size * seq_length, 256]);
31///
32/// assert_shape!(x, [batch_size, seq_length, 256]);     // names from dims() and a literal
33/// assert_shape!(flat, [batch_size * seq_length, 256]); // arithmetic on names
34/// assert_shape!(patch, [3 * 2, 2]);                    // arithmetic on literals
35/// assert_shape!(mask, [batch_size, _]);                // skip an axis
36/// ```
37///
38/// With `..`, the same check works whatever the rank:
39///
40/// ```
41/// use burn_tensor::{Tensor, assert_shape};
42///
43/// fn check_features<const D: usize>(x: &Tensor<D>, d_model: usize) {
44///     assert_shape!(x, [.., d_model]);
45/// }
46///
47/// fn check_channels<const D: usize>(x: &Tensor<D>, channels: usize) {
48///     assert_shape!(x, [_, channels, ..]);
49/// }
50///
51/// let device = Default::default();
52/// check_features(&Tensor::<2>::zeros([5, 256], &device), 256);
53/// check_features(&Tensor::<4>::zeros([2, 3, 5, 256], &device), 256);
54/// check_channels(&Tensor::<3>::zeros([2, 16, 100], &device), 16);
55/// check_channels(&Tensor::<5>::zeros([2, 16, 4, 8, 8], &device), 16);
56/// ```
57///
58/// A mismatch panics with the call and the offending axis:
59///
60/// ```should_panic
61/// use burn_tensor::{Tensor, assert_shape};
62/// let x = Tensor::<2>::zeros([2, 3], &Default::default());
63/// // assert_shape!(x, [2, 99]): axis 1 expected 99, got 3 (dims [2, 3])
64/// assert_shape!(x, [2, 99]);
65/// ```
66///
67/// A pattern whose length differs from the rank does not compile:
68///
69/// ```compile_fail,E0308
70/// use burn_tensor::{Tensor, assert_shape};
71/// let x = Tensor::<2>::zeros([2, 3], &Default::default());
72/// assert_shape!(x, [2, 3, 1]);
73/// ```
74///
75/// Neither does a slot that is not a `usize`:
76///
77/// ```compile_fail,E0308
78/// use burn_tensor::{Tensor, assert_shape};
79/// let x = Tensor::<2>::zeros([2, 3], &Default::default());
80/// assert_shape!(x, [2i32, 3]);
81/// ```
82///
83/// Only a `Tensor` is accepted. A `Shape` is a type error rather than a weaker check:
84///
85/// ```compile_fail,E0308
86/// use burn_tensor::{Tensor, assert_shape};
87/// let x = Tensor::<3>::zeros([2, 3, 4], &Default::default());
88/// assert_shape!(x.shape(), [2, 3]);
89/// ```
90///
91/// A pattern may contain `..` only once:
92///
93/// ```compile_fail
94/// use burn_tensor::{Tensor, assert_shape};
95/// let x = Tensor::<3>::zeros([2, 3, 4], &Default::default());
96/// assert_shape!(x, [.., 3, ..]);
97/// ```
98#[macro_export]
99macro_rules! assert_shape {
100    ($($tt:tt)*) => {
101        $crate::__assert_shape!($crate, $($tt)*)
102    };
103}
104
105/// Same as [`assert_shape!`](crate::assert_shape), but compiled out unless `debug_assertions`
106/// is on.
107///
108/// Without `..` the rank check is a type check and stays in every build; with `..` the
109/// minimum-rank check is a runtime assertion and is compiled out with the axis checks. Like
110/// `debug_assert!`, neither the tensor expression nor the slot expressions are evaluated in
111/// release builds.
112///
113/// ```
114/// use burn_tensor::{Tensor, debug_assert_shape};
115///
116/// let device = Default::default();
117/// let hidden = Tensor::<3>::zeros([2, 5, 256], &device);
118/// debug_assert_shape!(hidden, [2, _, 256]);
119/// ```
120///
121/// With debug assertions on, a mismatch panics like [`assert_shape!`](crate::assert_shape):
122///
123/// ```should_panic
124/// use burn_tensor::{Tensor, debug_assert_shape};
125/// let x = Tensor::<2>::zeros([2, 3], &Default::default());
126/// debug_assert_shape!(x, [2, 99]);
127/// # if !cfg!(debug_assertions) { panic!("compiled out") }
128/// ```
129#[macro_export]
130macro_rules! debug_assert_shape {
131    ($($tt:tt)*) => {
132        $crate::__debug_assert_shape!($crate, $($tt)*)
133    };
134}
burn_tensor/shape_macros.rs

burn_tensor/
shape_macros.rs
