tenferro-einsum 0.1.0

Subscripts, contraction planning, traced/eager einsum APIs, extension runtime, and AD rule for tenferro.
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

use tenferro_tensor::{Error, Result, TensorBackend, TensorScalar, TypedTensor};

use crate::eager::eager_einsum_read_subscripts;
use crate::Subscripts;

/// Execute eager einsum over typed tensors and return a typed result.
///
pub(crate) fn typed_eager_einsum<T: TensorScalar>(
    ctx: &mut impl TensorBackend,
    inputs: &[&TypedTensor<T>],
    subscripts: &str,
) -> Result<TypedTensor<T>> {
    let subscripts = Subscripts::parse(subscripts).map_err(|err| Error::InvalidConfig {
        op: "typed_eager_einsum",
        message: format!("invalid subscripts: {err}"),
    })?;
    let reads: Vec<_> = inputs.iter().map(|tensor| T::tensor_read(tensor)).collect();
    let result = eager_einsum_read_subscripts(ctx, &reads, &subscripts)?;
    let actual = result.dtype();
    T::into_typed(result).map_err(|_| Error::DTypeMismatch {
        op: "typed_eager_einsum",
        lhs: T::dtype(),
        rhs: actual,
    })
}