burn_tensor/tensor/api/einsum/
mod.rs

1//! Einstein summation through existing differentiable tensor operations.
2
3use core::marker::PhantomData;
4
5use crate::{Float, Tensor, kind::Numeric, ops::BridgeTensor};
6
7mod execution;
8
9/// An einsum operand whose rank is determined at runtime.
10///
11/// Allows [`Tensor::einsum`] to accept tensors of different ranks.
12/// All operands must have the same tensor kind, dtype and device.
13#[derive(Clone, Debug)]
14pub struct EinsumOperand<K: Numeric = Float> {
15    primitive: BridgeTensor,
16    kind: PhantomData<K>,
17}
18
19impl<const D: usize, K: Numeric> From<Tensor<D, K>> for EinsumOperand<K> {
20    fn from(tensor: Tensor<D, K>) -> Self {
21        Self {
22            primitive: tensor.primitive,
23            kind: PhantomData,
24        }
25    }
26}
27
28impl<const D: usize, K: Numeric> From<&Tensor<D, K>> for EinsumOperand<K> {
29    fn from(tensor: &Tensor<D, K>) -> Self {
30        tensor.clone().into()
31    }
32}
33
34impl<const D: usize, K: Numeric> Tensor<D, K> {
35    /// Computes Einstein summation with a runtime equation.
36    ///
37    /// Supports explicit (`"ij,jk->ik"`) and implicit (`"ij,jk"`) outputs,
38    /// any number of operands, repeated-label diagonals, and broadcasting.
39    /// See [`einsum!`](crate::einsum) for equation syntax and broadcasting rules.
40    /// The macro also checks literal equations and statically determined ranks
41    /// at compile time.
42    ///
43    /// Operands are contracted from left to right using multiplication,
44    /// reductions, and batched matrix multiplication. No optimized contraction
45    /// path is searched. Float operations retain automatic differentiation.
46    ///
47    /// Burn represents scalars as rank-one tensors with shape `[1]`: a scalar
48    /// result requires `D = 1`, and an empty input subscript accepts shape `[1]`.
49    /// Other inputs must have exactly the rank described by their subscripts.
50    ///
51    /// # Example
52    ///
53    /// ```
54    /// use burn_tensor::Tensor;
55    /// let device = Default::default();
56    /// let matrix = Tensor::<2>::from_floats([[1., 2.], [3., 4.]], &device);
57    /// let vector = Tensor::<1>::from_floats([10., 20.], &device);
58    /// let result = Tensor::<1>::einsum("ij,j->i", [matrix.into(), vector.into()]);
59    /// assert_eq!(result.into_data().try_to_vec::<f32>().unwrap(), [50., 110.]);
60    /// ```
61    ///
62    /// # Panics
63    ///
64    /// Panics for an invalid equation, no operands, incompatible ranks or
65    /// dimensions, unequal repeated-label dimensions, mismatched devices or
66    /// dtypes, quantized operands, or an incorrect output rank `D`.
67    pub fn einsum(equation: &str, operands: impl IntoIterator<Item = EinsumOperand<K>>) -> Self {
68        let equation =
69            burn_einsum::parse(equation).unwrap_or_else(|error| panic!("einsum: {error}"));
70        execution::execute(&equation.plan(), operands)
71    }
72}
73
74/// Implementation details used by the generated `einsum!` operation chain.
75#[doc(hidden)]
76pub mod __einsum {
77    pub use super::execution::{
78        alignment_shape, axes, broadcast_contract, can_contract_early, can_matmul, contract_early,
79        matmul_shapes, prepare, validate_broadcast,
80    };
81    pub use burn_einsum::Axis;
82}
burn_tensor/tensor/api/einsum/mod.rs

burn_tensor/tensor/api/einsum/
mod.rs
