burn-tensor 0.22.0

Tensor library with user-friendly APIs and automatic differentiation support
Documentation 
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

//! Einstein summation through existing differentiable tensor operations.

use core::marker::PhantomData;

use crate::{Float, Tensor, kind::Numeric, ops::BridgeTensor};

mod execution;

/// An einsum operand whose rank is determined at runtime.
///
/// Allows [`Tensor::einsum`] to accept tensors of different ranks.
/// All operands must have the same tensor kind, dtype and device.
#[derive(Clone, Debug)]
pub struct EinsumOperand<K: Numeric = Float> {
    primitive: BridgeTensor,
    kind: PhantomData<K>,
}

impl<const D: usize, K: Numeric> From<Tensor<D, K>> for EinsumOperand<K> {
    fn from(tensor: Tensor<D, K>) -> Self {
        Self {
            primitive: tensor.primitive,
            kind: PhantomData,
        }
    }
}

impl<const D: usize, K: Numeric> From<&Tensor<D, K>> for EinsumOperand<K> {
    fn from(tensor: &Tensor<D, K>) -> Self {
        tensor.clone().into()
    }
}

impl<const D: usize, K: Numeric> Tensor<D, K> {
    /// Computes Einstein summation with a runtime equation.
    ///
    /// Supports explicit (`"ij,jk->ik"`) and implicit (`"ij,jk"`) outputs,
    /// any number of operands, repeated-label diagonals, and broadcasting.
    /// See [`einsum!`](crate::einsum) for equation syntax and broadcasting rules.
    /// The macro also checks literal equations and statically determined ranks
    /// at compile time.
    ///
    /// Operands are contracted from left to right using multiplication,
    /// reductions, and batched matrix multiplication. No optimized contraction
    /// path is searched. Float operations retain automatic differentiation.
    ///
    /// Burn represents scalars as rank-one tensors with shape `[1]`: a scalar
    /// result requires `D = 1`, and an empty input subscript accepts shape `[1]`.
    /// Other inputs must have exactly the rank described by their subscripts.
    ///
    /// # Example
    ///
    /// ```
    /// use burn_tensor::Tensor;
    /// let device = Default::default();
    /// let matrix = Tensor::<2>::from_floats([[1., 2.], [3., 4.]], &device);
    /// let vector = Tensor::<1>::from_floats([10., 20.], &device);
    /// let result = Tensor::<1>::einsum("ij,j->i", [matrix.into(), vector.into()]);
    /// assert_eq!(result.into_data().try_to_vec::<f32>().unwrap(), [50., 110.]);
    /// ```
    ///
    /// # Panics
    ///
    /// Panics for an invalid equation, no operands, incompatible ranks or
    /// dimensions, unequal repeated-label dimensions, mismatched devices or
    /// dtypes, quantized operands, or an incorrect output rank `D`.
    pub fn einsum(equation: &str, operands: impl IntoIterator<Item = EinsumOperand<K>>) -> Self {
        let equation =
            burn_einsum::parse(equation).unwrap_or_else(|error| panic!("einsum: {error}"));
        execution::execute(&equation.plan(), operands)
    }
}

/// Implementation details used by the generated `einsum!` operation chain.
#[doc(hidden)]
pub mod __einsum {
    pub use super::execution::{
        alignment_shape, axes, broadcast_contract, can_contract_early, can_matmul, contract_early,
        matmul_shapes, prepare, validate_broadcast,
    };
    pub use burn_einsum::Axis;
}