burn_tensor/
einsum_macros.rs

1//! Literal Einstein summation with equation planning during compilation.
2
3/// Computes Einstein summation from a literal equation.
4///
5/// Give each input axis a letter, separate operands with commas, and use `->`
6/// to specify the output axes and their order. Values are multiplied along
7/// matching labels and summed over labels omitted from the output. For example,
8/// `"ij,jk->ik"` multiplies two matrices and sums over their shared `j` axis.
9///
10/// | Equation | Operation |
11/// | --- | --- |
12/// | `"ij,jk->ik"` | Matrix multiplication |
13/// | `"ij->ji"` | Transpose |
14/// | `"ii->i"` | Extract a diagonal |
15/// | `"ii->"` | Trace (sum of the diagonal) |
16/// | `"i,i->"` | Dot product |
17/// | `"...ij,...jk->...ik"` | Matrix multiplication with broadcast batch dimensions |
18///
19/// For dynamic equation strings, use [`Tensor::einsum`](crate::Tensor::einsum).
20///
21/// # Examples
22///
23/// ```
24/// use burn_tensor::{Tensor, einsum};
25/// let device = Default::default();
26/// let a = Tensor::<2>::from_floats([[1., 2.], [3., 4.]], &device);
27/// let b = Tensor::<2>::from_floats([[5., 6.], [7., 8.]], &device);
28/// let c = einsum!("ij,jk->ik", &a, b);
29/// assert_eq!(c.into_data().try_to_vec::<f32>().unwrap(), [19., 22., 43., 50.]);
30///
31/// let transposed = einsum!("ij->ji", &a);
32/// assert_eq!(transposed.into_data().try_to_vec::<f32>().unwrap(), [1., 3., 2., 4.]);
33/// let trace = einsum!("ii->", a);
34/// assert_eq!(trace.dims(), [1]);
35/// assert_eq!(trace.into_data().try_to_vec::<f32>().unwrap(), [5.]);
36/// ```
37///
38/// # Broadcasting and equation rules
39///
40/// Matching labels across operands must have equal sizes or a size of one, which
41/// broadcasts to the other size. Repeating a label within one operand extracts
42/// its diagonal: those axis sizes must be equal, even if one is a singleton.
43///
44/// `...` matches zero or more axes. Ellipsis dimensions broadcast from the right,
45/// allowing operands to have different numbers of batch dimensions. To sum over
46/// these dimensions, omit `...` from the explicit output.
47///
48/// ```
49/// use burn_tensor::{Tensor, einsum};
50/// let device = Default::default();
51/// let a = Tensor::<3>::ones([2, 3, 4], &device);
52/// let b = Tensor::<2>::ones([4, 5], &device);
53/// // b is shared across both batches of a.
54/// let c: Tensor<3> = einsum!("...ij,...jk->...ik", a, b);
55/// assert_eq!(c.dims(), [2, 3, 5]);
56/// ```
57///
58/// Each axis label is a single letter from `a-z` or `A-Z`. Uppercase and lowercase
59/// letters are distinct labels. Spaces between tokens are ignored.
60/// With no `->`, the output contains the ellipsis first, followed by labels
61/// occurring exactly once across all inputs, sorted `A-Z`, then `a-z`.
62/// For example, `"ij,jk"` is equivalent to `"ij,jk->ik"`.
63///
64/// # Types and limitations
65///
66/// Float and Int operands must share their kind, dtype, and device. Float
67/// operations support automatic differentiation; quantized operands are unsupported.
68/// Scalar results have shape `[1]`, and an empty input subscript accepts shape `[1]`.
69/// With an output ellipsis of unknown width, supply the output type, as above.
70/// Operands are contracted from left to right; no optimized contraction order is searched.
71///
72/// # Validation
73///
74/// Equations, operand counts, and exact ranks are checked at compile time where
75/// the equation determines them. Axis sizes and ellipsis widths are checked at runtime.
76///
77/// Unknown output labels are compile errors:
78///
79/// ```compile_fail
80/// use burn_tensor::{Tensor, einsum};
81/// let a = Tensor::<1>::zeros([3], &Default::default());
82/// let _ = einsum!("i->j", a);
83/// ```
84///
85/// So are incorrect input ranks:
86///
87/// ```compile_fail
88/// use burn_tensor::{Tensor, einsum};
89/// let a = Tensor::<1>::zeros([3], &Default::default());
90/// let _ = einsum!("ij->ji", a);
91/// ```
92///
93/// And incorrect output ranks:
94///
95/// ```compile_fail
96/// use burn_tensor::{Tensor, einsum};
97/// let a = Tensor::<2>::zeros([2, 3], &Default::default());
98/// let _: Tensor<1> = einsum!("ij->ji", a);
99/// ```
100/// An equation must have one subscript per operand:
101///
102/// ```compile_fail
103/// use burn_tensor::{Tensor, einsum};
104/// let a = Tensor::<1>::zeros([3], &Default::default());
105/// let _ = einsum!("i,j->ij", a);
106/// ```
107///
108/// Malformed ellipses are rejected during compilation:
109///
110/// ```compile_fail
111/// use burn_tensor::{Tensor, einsum};
112/// let a = Tensor::<1>::zeros([3], &Default::default());
113/// let _ = einsum!("i..->i", a);
114/// ```
115///
116/// # Panics
117///
118/// Panics for incompatible ranks or broadcast dimensions, unequal repeated-label
119/// dimensions, mismatched devices or dtypes, quantized operands, or an incorrect
120/// output rank when the output contains an ellipsis of unknown width.
121#[macro_export]
122macro_rules! einsum {
123    ($($tt:tt)*) => {
124        $crate::__einsum!($crate, $($tt)*)
125    };
126}
burn_tensor/einsum_macros.rs

burn_tensor/
einsum_macros.rs
