tenferro-einsum 0.3.0

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

use crate::{Result, Subscripts};

/// Canonical N-ary einsum subscripts using integer labels.
///
/// String notation is a user-facing convenience. Runtime integration layers can
/// carry this representation in extension payloads so execution, shape
/// inference, and AD do not need to parse strings.
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub struct EinsumSubscripts {
    /// Index labels for each input tensor.
    pub inputs: Vec<Vec<u32>>,
    /// Index labels for the output tensor.
    pub output: Vec<u32>,
}

impl EinsumSubscripts {
    /// Create subscripts from integer label arrays.
    ///
    /// # Examples
    ///
    /// ```
    /// use tenferro_einsum::EinsumSubscripts;
    ///
    /// let subscripts = EinsumSubscripts::new(&[&[0, 1], &[1, 2]], &[0, 2]);
    ///
    /// assert_eq!(subscripts.inputs, vec![vec![0, 1], vec![1, 2]]);
    /// assert_eq!(subscripts.output, vec![0, 2]);
    /// ```
    pub fn new(inputs: &[&[u32]], output: &[u32]) -> Self {
        Self {
            inputs: inputs.iter().map(|labels| labels.to_vec()).collect(),
            output: output.to_vec(),
        }
    }

    /// Number of input operands described by this specification.
    ///
    /// # Examples
    ///
    /// ```
    /// use tenferro_einsum::EinsumSubscripts;
    ///
    /// let subscripts = EinsumSubscripts::new(&[&[0], &[0]], &[]);
    ///
    /// assert_eq!(subscripts.input_count(), 2);
    /// ```
    #[must_use]
    pub fn input_count(&self) -> usize {
        self.inputs.len()
    }
}

impl From<Subscripts> for EinsumSubscripts {
    fn from(subscripts: Subscripts) -> Self {
        Self {
            inputs: subscripts.inputs,
            output: subscripts.output,
        }
    }
}

impl From<&Subscripts> for EinsumSubscripts {
    fn from(subscripts: &Subscripts) -> Self {
        Self {
            inputs: subscripts.inputs.clone(),
            output: subscripts.output.clone(),
        }
    }
}

impl From<EinsumSubscripts> for Subscripts {
    fn from(subscripts: EinsumSubscripts) -> Self {
        Self {
            inputs: subscripts.inputs,
            output: subscripts.output,
        }
    }
}

impl From<&EinsumSubscripts> for Subscripts {
    fn from(subscripts: &EinsumSubscripts) -> Self {
        Self {
            inputs: subscripts.inputs.clone(),
            output: subscripts.output.clone(),
        }
    }
}

/// Parse string einsum notation into canonical integer labels.
///
/// # Examples
///
/// ```
/// use tenferro_einsum::parse_einsum_subscripts;
///
/// let subscripts = parse_einsum_subscripts("ij,jk->ik").unwrap();
///
/// assert_eq!(subscripts.inputs.len(), 2);
/// assert_eq!(subscripts.output, vec![b'i' as u32, b'k' as u32]);
/// ```
///
/// # Errors
///
/// Returns [`crate::Error::InvalidSubscripts`] when the notation is malformed,
/// contains an invalid label, or has an invalid input/output separator.
pub fn parse_einsum_subscripts(notation: &str) -> Result<EinsumSubscripts> {
    Subscripts::parse(notation).map(EinsumSubscripts::from)
}

#[cfg(test)]
mod tests;