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
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
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125

use crate::syntax::subscripts::Subscripts;
use crate::{Error, Result as EinsumResult};

use super::{compile_pairwise_step_plan, compile_strict_binary_lowering_step_plan};

fn compile_strict_step_plan(
    subs: &Subscripts,
    shapes: &[&[usize]],
) -> EinsumResult<Option<super::StrictBinaryLoweringPlan>> {
    let size_dict = crate::util::build_size_dict(subs, shapes, None)?;
    compile_strict_binary_lowering_step_plan(
        &subs.inputs[0],
        &subs.inputs[1],
        &subs.output,
        &size_dict,
    )
}

#[test]
fn strict_binary_lowering_plan_builds_dense_matmul() {
    let subs = Subscripts::new(&[&[0, 1], &[1, 2]], &[0, 2]);
    let shapes = [&[2, 3][..], &[3, 4][..]];

    let plan = compile_strict_step_plan(&subs, &shapes)
        .unwrap()
        .expect("simple dense matmul should admit strict lowering");

    assert_eq!(plan.lhs_free_labels, vec![0]);
    assert_eq!(plan.rhs_free_labels, vec![2]);
    assert_eq!(plan.contract_labels, vec![1]);
    assert_eq!(plan.lhs_perm, vec![0, 1]);
    assert_eq!(plan.rhs_perm, vec![0, 1]);
    assert_eq!(plan.m, 2);
    assert_eq!(plan.k, 3);
    assert_eq!(plan.n, 4);
    assert_eq!(plan.lhs_matrix_dims, vec![2, 3]);
    assert_eq!(plan.rhs_matrix_dims, vec![3, 4]);
    assert_eq!(plan.canonical_output_dims, vec![2, 4]);
    assert_eq!(plan.output_perm, vec![0, 1]);
}

#[test]
fn strict_binary_lowering_rejects_repeated_labels() {
    let subs = Subscripts::new(&[&[0, 0, 1], &[1, 2]], &[0, 2]);
    let shapes = [&[2, 2, 3][..], &[3, 4][..]];

    let plan = compile_strict_step_plan(&subs, &shapes).unwrap();

    assert!(
        plan.is_none(),
        "repeated labels must stay on the generic binary lowering path"
    );
}

#[test]
fn strict_binary_lowering_plan_tracks_non_identity_output_permutation() {
    let subs = Subscripts::new(&[&[0, 1], &[1, 2]], &[2, 0]);
    let shapes = [&[2, 3][..], &[3, 4][..]];

    let plan = compile_strict_step_plan(&subs, &shapes)
        .unwrap()
        .expect("output permutations should still admit strict lowering");

    assert_eq!(plan.canonical_output_dims, vec![2, 4]);
    assert_eq!(plan.output_perm, vec![1, 0]);
}

#[test]
fn pairwise_step_plan_embeds_strict_binary_lowering_when_available() {
    let subs = Subscripts::new(&[&[0, 1], &[1, 2]], &[2, 0]);
    let shapes = [&[2, 3][..], &[3, 4][..]];
    let size_dict = crate::util::build_size_dict(&subs, &shapes, None).unwrap();

    let step_plan =
        compile_pairwise_step_plan(&subs.inputs[0], &subs.inputs[1], &subs.output, &size_dict)
            .unwrap();

    let strict = step_plan
        .strict_binary
        .as_ref()
        .expect("dense binary matmul step should cache strict lowering");
    assert_eq!(strict.output_perm, vec![1, 0]);
}

#[test]
fn pairwise_step_plan_preserves_strict_lowering_error_type() {
    let size_dict = [(0, 2), (1, 3)].into_iter().collect();

    let err = compile_pairwise_step_plan(&[0], &[1], &[2], &size_dict).unwrap_err();

    assert!(matches!(
        err,
        Error::InvalidArgument(message)
            if message.contains("unknown size")
    ));
}

#[test]
fn pairwise_step_plan_rejects_overflowing_fused_dimensions() {
    let size_dict = [(0, usize::MAX), (1, 2), (2, 3), (4, 2)]
        .into_iter()
        .collect();

    let err = compile_pairwise_step_plan(&[0, 4, 1], &[1, 2], &[0, 4, 2], &size_dict).unwrap_err();

    assert!(matches!(
        err,
        Error::InvalidArgument(message)
            if message.contains("dimension product overflow")
    ));
}

#[test]
fn strict_binary_lowering_rejects_overflowing_fused_dimensions() {
    let subs = Subscripts::new(&[&[0, 4, 1], &[1, 2]], &[0, 4, 2]);
    let shapes = [&[usize::MAX, 2, 2][..], &[2, 3][..]];

    let err = compile_strict_step_plan(&subs, &shapes).unwrap_err();

    assert!(matches!(
        err,
        Error::InvalidArgument(message)
            if message.contains("dimension product overflow")
    ));
}