tenferro_einsum/
lib.rs

1//! High-level einsum with N-ary contraction tree optimization.
2//!
3//! This crate provides:
4//!
5//! - **String notation**: `"ij,jk->ik"` (NumPy/PyTorch compatible)
6//! - **Parenthesized notation**: `"ij,(jk,kl)->il"` respects user-specified
7//!   contraction order via [`NestedEinsum`]
8//! - **Integer label notation**: using `u32` labels
9//! - **Repeated labels**: `"ii->i"` extracts diagonals, `"ii->"` traces, and
10//!   `"i->ii"` embeds a vector on a diagonal
11//! - **N-ary contraction**: Automatic or manual optimization of pairwise
12//!   contraction order via [`ContractionTree`]
13//! - **Tensordot sugar**: NumPy-style axis-pair contraction extension methods,
14//!   implemented as contraction sugar rather than as linear algebra APIs.
15//! - **Concrete execution**: backend-explicit [`TensorEinsumExt`],
16//!   [`TypedTensorEinsumExt`], [`TensorReadEinsumExt`], and
17//!   [`ConcreteEinsumPlan`] APIs for non-AD tensor values.
18//! - **Extension runtime**: traced einsum lowers to a registered tenferro
19//!   extension runtime, keeping core op definitions small.
20//! - **Tensor extension traits**: graph-building and immediate-execution
21//!   helpers are available as methods on `GraphCompiler`, concrete input
22//!   slices/arrays, eager input slices/arrays, and tensor receivers.
23//!
24//! # Examples
25//!
26//! ```
27//! use tenferro_einsum::{ContractionTree, Subscripts};
28//!
29//! let subs = Subscripts::parse("ij,jk->ik").unwrap();
30//! let tree = ContractionTree::optimize(&subs, &[&[2, 3], &[3, 4]]).unwrap();
31//! assert_eq!(tree.step_count(), 1);
32//! ```
33//!
34//! ```
35//! use tenferro_cpu::CpuBackend;
36//! use tenferro_einsum::TensorEinsumExt;
37//! use tenferro_tensor::Tensor;
38//!
39//! let a = Tensor::from_vec_col_major(vec![2, 3], vec![1.0_f64; 6]).unwrap();
40//! let b = Tensor::from_vec_col_major(vec![3, 4], vec![1.0_f64; 12]).unwrap();
41//! let mut backend = CpuBackend::new();
42//!
43//! let out = [&a, &b].einsum("ij,jk->ik", &mut backend)?;
44//! assert_eq!(out.shape(), &[2, 4]);
45//! # Ok::<(), tenferro_tensor::Error>(())
46//! ```
47//!
48//! ```
49//! use tenferro_einsum::Subscripts;
50//!
51//! let trace = Subscripts::parse("ii->").unwrap();
52//! let diagonal = Subscripts::parse("ii->i").unwrap();
53//! let embedded = Subscripts::parse("i->ii").unwrap();
54//! let higher_rank = Subscripts::parse("iij->ij").unwrap();
55//!
56//! assert!(trace.output.is_empty());
57//! assert_eq!(diagonal.output, vec![b'i' as u32]);
58//! assert_eq!(embedded.output, vec![b'i' as u32, b'i' as u32]);
59//! assert_eq!(higher_rank.inputs[0], vec![b'i' as u32, b'i' as u32, b'j' as u32]);
60//! ```
61
62mod binary_dot;
63mod builder;
64mod cache;
65mod concrete;
66mod eager;
67#[cfg(feature = "autodiff")]
68mod eager_ad;
69mod error;
70mod extension;
71pub mod lowering;
72mod optimize;
73mod planning;
74mod subscripts;
75mod syntax;
76mod tensordot;
77mod traced;
78#[cfg(test)]
79mod typed_eager;
80pub(crate) mod util;
81
82pub use cache::EINSUM_EXTENSION_FAMILY_ID;
83pub use concrete::{
84    ConcreteEinsumPlan, TensorEinsumExt, TensorEinsumIntoExt, TensorReadEinsumExt,
85    TensorReadEinsumIntoExt, TypedTensorEinsumExt, TypedTensorEinsumIntoExt,
86};
87#[cfg(feature = "autodiff")]
88pub use eager_ad::{EagerEinsumExt, EagerTensorEinsumExt};
89pub use error::{Error, Result};
90#[cfg(feature = "autodiff")]
91pub use extension::ad_rules;
92pub use extension::register_runtime;
93pub use optimize::EinsumOptimize;
94pub use planning::tree::{ContractionOptimizerOptions, ContractionTree};
95pub use subscripts::{parse_einsum_subscripts, EinsumSubscripts};
96pub use syntax::nested::NestedEinsum;
97pub use syntax::subscripts::Subscripts;
98pub use tensordot::TensorDotAxes;
99pub use traced::{GraphCompilerEinsumExt, TracedTensorEinsumExt};
100
101#[cfg(test)]
102mod concrete_tests;
103#[cfg(test)]
104mod tests;
105#[cfg(test)]
106mod typed_eager_tests;
tenferro_einsum/lib.rs

tenferro_einsum/
lib.rs
