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
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
//! Shape assertion macros. Parsing and expansion live in burn-derive; these wrappers forward
//! `$crate` so the expansion can name [`Tensor::dims`](crate::Tensor::dims) by path, which
//! makes anything other than a `Tensor` a type error.
//!
//! Inspired by the `burn-contracts` crate by Crutcher Dunnavant.
/// Asserts a tensor's rank, at compile time unless the pattern has `..`, and its axis sizes at
/// runtime.
///
/// The pattern has one slot per axis. A slot is either `_`, which skips the axis, or any `usize`
/// expression the axis must equal: a name from `dims()`, a config field, a literal, arithmetic
/// on any of those. To name axes, destructure `tensor.dims()` first and use the names as slots.
///
/// The pattern length must equal the tensor rank, otherwise the call does not compile. One
/// slot may be `..`, which stands for any number of axes; the pattern then fixes only a
/// minimum rank, checked at runtime, and fits a function generic over `const D: usize`. A
/// pattern of just `[..]` checks nothing beyond the argument being a tensor. The tensor is
/// borrowed, not moved, and every expression is evaluated once. Checks stay on in release
/// builds; see [`debug_assert_shape!`](crate::debug_assert_shape) for the debug-only variant.
///
/// ```
/// use burn_tensor::{Tensor, assert_shape};
///
/// let device = Default::default();
/// let x = Tensor::<3>::zeros([2, 5, 256], &device);
/// let mask = Tensor::<2>::zeros([2, 5], &device);
/// let patch = Tensor::<2>::zeros([6, 2], &device);
///
/// let [batch_size, seq_length, _] = x.dims();
/// let flat = x.clone().reshape([batch_size * seq_length, 256]);
///
/// assert_shape!(x, [batch_size, seq_length, 256]); // names from dims() and a literal
/// assert_shape!(flat, [batch_size * seq_length, 256]); // arithmetic on names
/// assert_shape!(patch, [3 * 2, 2]); // arithmetic on literals
/// assert_shape!(mask, [batch_size, _]); // skip an axis
/// ```
///
/// With `..`, the same check works whatever the rank:
///
/// ```
/// use burn_tensor::{Tensor, assert_shape};
///
/// fn check_features<const D: usize>(x: &Tensor<D>, d_model: usize) {
/// assert_shape!(x, [.., d_model]);
/// }
///
/// fn check_channels<const D: usize>(x: &Tensor<D>, channels: usize) {
/// assert_shape!(x, [_, channels, ..]);
/// }
///
/// let device = Default::default();
/// check_features(&Tensor::<2>::zeros([5, 256], &device), 256);
/// check_features(&Tensor::<4>::zeros([2, 3, 5, 256], &device), 256);
/// check_channels(&Tensor::<3>::zeros([2, 16, 100], &device), 16);
/// check_channels(&Tensor::<5>::zeros([2, 16, 4, 8, 8], &device), 16);
/// ```
///
/// A mismatch panics with the call and the offending axis:
///
/// ```should_panic
/// use burn_tensor::{Tensor, assert_shape};
/// let x = Tensor::<2>::zeros([2, 3], &Default::default());
/// // assert_shape!(x, [2, 99]): axis 1 expected 99, got 3 (dims [2, 3])
/// assert_shape!(x, [2, 99]);
/// ```
///
/// A pattern whose length differs from the rank does not compile:
///
/// ```compile_fail,E0308
/// use burn_tensor::{Tensor, assert_shape};
/// let x = Tensor::<2>::zeros([2, 3], &Default::default());
/// assert_shape!(x, [2, 3, 1]);
/// ```
///
/// Neither does a slot that is not a `usize`:
///
/// ```compile_fail,E0308
/// use burn_tensor::{Tensor, assert_shape};
/// let x = Tensor::<2>::zeros([2, 3], &Default::default());
/// assert_shape!(x, [2i32, 3]);
/// ```
///
/// Only a `Tensor` is accepted. A `Shape` is a type error rather than a weaker check:
///
/// ```compile_fail,E0308
/// use burn_tensor::{Tensor, assert_shape};
/// let x = Tensor::<3>::zeros([2, 3, 4], &Default::default());
/// assert_shape!(x.shape(), [2, 3]);
/// ```
///
/// A pattern may contain `..` only once:
///
/// ```compile_fail
/// use burn_tensor::{Tensor, assert_shape};
/// let x = Tensor::<3>::zeros([2, 3, 4], &Default::default());
/// assert_shape!(x, [.., 3, ..]);
/// ```
/// Same as [`assert_shape!`](crate::assert_shape), but compiled out unless `debug_assertions`
/// is on.
///
/// Without `..` the rank check is a type check and stays in every build; with `..` the
/// minimum-rank check is a runtime assertion and is compiled out with the axis checks. Like
/// `debug_assert!`, neither the tensor expression nor the slot expressions are evaluated in
/// release builds.
///
/// ```
/// use burn_tensor::{Tensor, debug_assert_shape};
///
/// let device = Default::default();
/// let hidden = Tensor::<3>::zeros([2, 5, 256], &device);
/// debug_assert_shape!(hidden, [2, _, 256]);
/// ```
///
/// With debug assertions on, a mismatch panics like [`assert_shape!`](crate::assert_shape):
///
/// ```should_panic
/// use burn_tensor::{Tensor, debug_assert_shape};
/// let x = Tensor::<2>::zeros([2, 3], &Default::default());
/// debug_assert_shape!(x, [2, 99]);
/// # if !cfg!(debug_assertions) { panic!("compiled out") }
/// ```