tenferro-einsum 0.2.0

Subscripts, contraction planning, concrete/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
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

use tenferro_ops::std_tensor_op::StdTensorOp;
use tenferro_runtime::{DType, GraphCompiler, TracedTensor};

use super::{einsum, einsum_with};
use crate::EinsumOptimize;

#[test]
fn concrete_traced_nary_einsum_expands_to_standard_graph() {
    let a = TracedTensor::from_vec_col_major(vec![2, 3], vec![1.0_f64; 6]).unwrap();
    let b = TracedTensor::from_vec_col_major(vec![3, 4], vec![1.0_f64; 12]).unwrap();
    let c = TracedTensor::from_vec_col_major(vec![4, 5], vec![1.0_f64; 20]).unwrap();
    let mut compiler = GraphCompiler::new();

    let out = einsum(&mut compiler, &[&a, &b, &c], "ij,jk,kl->il").unwrap();

    assert!(out
        .graph()
        .operations()
        .iter()
        .all(|node| { !matches!(node.operation, StdTensorOp::Extension(_)) }));
    assert!(out
        .graph()
        .operations()
        .iter()
        .any(|node| { matches!(node.operation, StdTensorOp::DotGeneral { .. }) }));
}

#[test]
fn symbolic_path_traced_nary_einsum_expands_to_standard_graph() {
    let a = TracedTensor::input_symbolic_shape(DType::F64, 2).unwrap();
    let b = TracedTensor::input_symbolic_shape(DType::F64, 2).unwrap();
    let c = TracedTensor::input_symbolic_shape(DType::F64, 2).unwrap();
    let mut compiler = GraphCompiler::new();

    let out = einsum_with(
        &mut compiler,
        &[&a, &b, &c],
        "ij,jk,kl->il",
        EinsumOptimize::Path(vec![(0, 1), (0, 1)]),
    )
    .unwrap();

    assert!(out
        .graph()
        .operations()
        .iter()
        .all(|node| { !matches!(node.operation, StdTensorOp::Extension(_)) }));
    assert!(out
        .graph()
        .operations()
        .iter()
        .any(|node| { matches!(node.operation, StdTensorOp::DotGeneral { .. }) }));
}

#[test]
fn symbolic_auto_traced_nary_einsum_remains_extension() {
    let a = TracedTensor::input_symbolic_shape(DType::F64, 2).unwrap();
    let b = TracedTensor::input_symbolic_shape(DType::F64, 2).unwrap();
    let c = TracedTensor::input_symbolic_shape(DType::F64, 2).unwrap();
    let mut compiler = GraphCompiler::new();

    let out = einsum(&mut compiler, &[&a, &b, &c], "ij,jk,kl->il").unwrap();

    assert!(out
        .graph()
        .operations()
        .iter()
        .any(|node| { matches!(node.operation, StdTensorOp::Extension(_)) }));
}