einsum_ndarray/
lib.rs

1#![forbid(unsafe_code)]
2#![warn(missing_docs)]
3
4//! Einstein summation for dynamically shaped `ndarray` arrays.
5//!
6//! Integer additions and multiplications use wrapping arithmetic in every
7//! build profile. Other element types use their `ndarray` arithmetic.
8
9mod execute;
10mod parse;
11mod plan;
12
13use std::error::Error;
14use std::fmt;
15
16use ndarray::{ArrayD, ArrayViewD, LinalgScalar};
17
18use parse::ParsedExpression;
19
20/// Selects how a contraction path is built.
21#[derive(Clone, Debug, Default, Eq, PartialEq)]
22#[non_exhaustive]
23pub enum Strategy {
24    /// Searches all contraction trees for up to eight operands and uses a
25    /// greedy search above that limit.
26    #[default]
27    Auto,
28    /// Searches all contraction trees without the eight-operand cutoff.
29    ///
30    /// Expressions above eight operands currently use the greedy search to
31    /// keep planning time bounded.
32    Optimal,
33    /// Repeatedly contracts the pair with the lowest greedy score.
34    Greedy,
35    /// Uses caller-supplied current-list indices.
36    ///
37    /// Each step removes its operands and appends its intermediate.
38    Explicit(Vec<Vec<usize>>),
39}
40
41/// Describes an ordered contraction.
42#[derive(Clone, Debug, Eq, PartialEq)]
43pub struct ContractionPath {
44    /// Steps in execution order.
45    pub steps: Vec<ContractionStep>,
46    /// Cost of evaluating every operand in one expression.
47    pub naive_flops: u128,
48    /// Sum of the selected step costs.
49    pub optimized_flops: u128,
50    /// Largest intermediate element count.
51    pub largest_intermediate: u128,
52}
53
54/// Describes one contraction step.
55#[derive(Clone, Debug, Eq, PartialEq)]
56pub struct ContractionStep {
57    /// Indices into the operand list as it exists before this step.
58    pub operands: Vec<usize>,
59    /// Fully resolved input and output labels for this step.
60    pub subscripts: String,
61    /// Cost assigned by the contraction model.
62    pub flops: u128,
63    /// Whether execution uses `ndarray` matrix multiplication for this step.
64    pub gemm: bool,
65}
66
67/// Reports invalid expressions, shapes, and paths.
68#[derive(Clone, Debug, Eq, PartialEq)]
69#[non_exhaustive]
70pub enum EinsumError {
71    /// A byte outside the subscript grammar was found.
72    InvalidCharacter {
73        /// Byte position in the supplied string.
74        position: usize,
75        /// Invalid byte.
76        byte: u8,
77    },
78    /// An arrow or another grammar element is malformed.
79    MalformedSubscripts {
80        /// Byte position in the supplied string.
81        position: usize,
82        /// Short explanation of the grammar error.
83        reason: &'static str,
84    },
85    /// The number of input groups differs from the operand count.
86    OperandCountMismatch {
87        /// Number of input groups.
88        subscripts: usize,
89        /// Number of operand shapes or arrays.
90        operands: usize,
91    },
92    /// An operand group names more axes than the operand has.
93    TooManyLabels {
94        /// Zero-based operand index.
95        operand: usize,
96        /// Number of explicit labels.
97        labels: usize,
98        /// Operand rank.
99        ndim: usize,
100    },
101    /// An operand without an ellipsis names fewer axes than its rank.
102    TooFewLabels {
103        /// Zero-based operand index.
104        operand: usize,
105        /// Number of labels.
106        labels: usize,
107        /// Operand rank.
108        ndim: usize,
109    },
110    /// Repeated axes in one operand have different lengths.
111    DiagonalSizeMismatch {
112        /// Zero-based operand index.
113        operand: usize,
114        /// Repeated label.
115        label: char,
116        /// First and conflicting lengths.
117        sizes: (usize, usize),
118    },
119    /// Two axes cannot broadcast to one label size.
120    BroadcastMismatch {
121        /// Named label, or `None` for an ellipsis axis.
122        label: Option<char>,
123        /// Conflicting lengths.
124        sizes: (usize, usize),
125    },
126    /// An explicit output repeats a label.
127    OutputLabelRepeated {
128        /// Repeated label.
129        label: char,
130    },
131    /// An explicit output names a label absent from every input.
132    OutputLabelUnknown {
133        /// Unknown label.
134        label: char,
135    },
136    /// An explicit output omits broadcast axes.
137    OutputEllipsisMissing {
138        /// Number of omitted broadcast axes.
139        broadcast_rank: usize,
140    },
141    /// A required array size does not fit in `usize`.
142    SizeOverflow,
143    /// Execution operands differ from the shapes bound to a plan.
144    ShapeMismatch {
145        /// Zero-based operand index.
146        operand: usize,
147        /// Shape stored in the plan.
148        expected: Vec<usize>,
149        /// Shape supplied for execution.
150        actual: Vec<usize>,
151    },
152    /// A supplied path cannot reduce the operand list to one value.
153    InvalidPath {
154        /// Short explanation of the path error.
155        reason: &'static str,
156    },
157}
158
159impl fmt::Display for EinsumError {
160    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
161        match self {
162            Self::InvalidCharacter { position, byte } => {
163                write!(formatter, "invalid byte {byte:#04x} at position {position}")
164            }
165            Self::MalformedSubscripts { position, reason } => {
166                write!(formatter, "malformed subscripts at position {position}: {reason}")
167            }
168            Self::OperandCountMismatch {
169                subscripts,
170                operands,
171            } => write!(
172                formatter,
173                "subscripts contain {subscripts} operands but {operands} were supplied"
174            ),
175            Self::TooManyLabels {
176                operand,
177                labels,
178                ndim,
179            } => write!(
180                formatter,
181                "operand {operand} has rank {ndim} but its subscript has {labels} labels"
182            ),
183            Self::TooFewLabels {
184                operand,
185                labels,
186                ndim,
187            } => write!(
188                formatter,
189                "operand {operand} has rank {ndim} but its subscript has {labels} labels and no ellipsis"
190            ),
191            Self::DiagonalSizeMismatch {
192                operand,
193                label,
194                sizes,
195            } => write!(
196                formatter,
197                "operand {operand} repeats label {label} on axes of lengths {} and {}",
198                sizes.0, sizes.1
199            ),
200            Self::BroadcastMismatch { label, sizes } => match label {
201                Some(label) => write!(
202                    formatter,
203                    "label {label} cannot broadcast lengths {} and {}",
204                    sizes.0, sizes.1
205                ),
206                None => write!(
207                    formatter,
208                    "ellipsis axis cannot broadcast lengths {} and {}",
209                    sizes.0, sizes.1
210                ),
211            },
212            Self::OutputLabelRepeated { label } => {
213                write!(formatter, "output label {label} appears more than once")
214            }
215            Self::OutputLabelUnknown { label } => {
216                write!(formatter, "output label {label} does not appear in an input")
217            }
218            Self::OutputEllipsisMissing { broadcast_rank } => write!(
219                formatter,
220                "output omits an ellipsis with rank {broadcast_rank}"
221            ),
222            Self::SizeOverflow => formatter.write_str("array size exceeds usize"),
223            Self::ShapeMismatch {
224                operand,
225                expected,
226                actual,
227            } => write!(
228                formatter,
229                "operand {operand} has shape {actual:?}, expected {expected:?}"
230            ),
231            Self::InvalidPath { reason } => write!(formatter, "invalid path: {reason}"),
232        }
233    }
234}
235
236impl Error for EinsumError {}
237
238/// A parsed contraction bound to fixed operand shapes.
239#[derive(Clone, Debug)]
240pub struct EinsumPlan {
241    expression: ParsedExpression,
242    path: ContractionPath,
243    execution: execute::ExecutionPlan,
244    output_subscripts: String,
245}
246
247impl EinsumPlan {
248    /// Parses and plans an expression for fixed operand shapes.
249    pub fn new(subscripts: &str, shapes: &[&[usize]]) -> Result<Self, EinsumError> {
250        Self::with_strategy(subscripts, shapes, Strategy::Auto)
251    }
252
253    /// Parses and plans an expression with a selected path strategy.
254    pub fn with_strategy(
255        subscripts: &str,
256        shapes: &[&[usize]],
257        strategy: Strategy,
258    ) -> Result<Self, EinsumError> {
259        let expression = parse::parse(subscripts, shapes)?;
260        let path = plan::build_path(&expression, strategy)?;
261        let execution = execute::build_execution_plan(&expression, &path)?;
262        let output_subscripts = parse::format_labels(&expression.output, expression.broadcast_rank);
263        Ok(Self {
264            expression,
265            path,
266            execution,
267            output_subscripts,
268        })
269    }
270
271    /// Executes the plan on arrays with the planned shapes.
272    pub fn execute<A: LinalgScalar>(
273        &self,
274        operands: &[ArrayViewD<'_, A>],
275    ) -> Result<ArrayD<A>, EinsumError> {
276        execute::execute(self, operands)
277    }
278
279    /// Returns the result shape.
280    pub fn output_shape(&self) -> &[usize] {
281        &self.expression.output_shape
282    }
283
284    /// Returns the selected contraction path.
285    pub fn path(&self) -> &ContractionPath {
286        &self.path
287    }
288
289    /// Returns the resolved result labels.
290    ///
291    /// A broadcast block is written as `...`.
292    pub fn output_subscripts(&self) -> &str {
293        &self.output_subscripts
294    }
295}
296
297/// Parses, plans, and evaluates one expression.
298pub fn einsum<A: LinalgScalar>(
299    subscripts: &str,
300    operands: &[ArrayViewD<'_, A>],
301) -> Result<ArrayD<A>, EinsumError> {
302    let shapes: Vec<&[usize]> = operands.iter().map(ArrayViewD::shape).collect();
303    EinsumPlan::new(subscripts, &shapes)?.execute(operands)
304}
305
306/// Parses and plans one expression from shapes alone.
307pub fn einsum_path(subscripts: &str, shapes: &[&[usize]]) -> Result<ContractionPath, EinsumError> {
308    Ok(EinsumPlan::new(subscripts, shapes)?.path)
309}
einsum_ndarray/lib.rs

einsum_ndarray/
lib.rs
