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//! - **Ellipsis notation**: `"...ij,...jk->...ik"` with right-aligned
7//!   equal-or-one broadcasting, plus programmatic [`EinsumNotation`]
8//! - **Parenthesized notation**: `"ij,(jk,kl)->il"` respects user-specified
9//!   contraction order via [`NestedEinsum`]
10//! - **Integer label notation**: using `u32` labels
11//! - **Repeated labels**: `"ii->i"` extracts diagonals, `"ii->"` traces, and
12//!   `"i->ii"` embeds a vector on a diagonal
13//! - **N-ary contraction**: Automatic or manual optimization of pairwise
14//!   contraction order via [`ContractionTree`]
15//! - **Tensordot sugar**: NumPy-style axis-pair contraction extension methods,
16//!   implemented as contraction sugar rather than as linear algebra APIs.
17//! - **Concrete execution**: backend-explicit [`TensorEinsumExt`],
18//!   [`TypedTensorEinsumExt`], [`TensorReadEinsumExt`],
19//!   [`TypedTensorReadEinsumExt`], and [`ConcreteEinsumPlan`] APIs for non-AD
20//!   tensor values.
21//! - **Extension runtime**: traced einsum lowers to a registered tenferro
22//!   extension runtime, keeping core op definitions small.
23//! - **Tensor extension traits**: graph-building and immediate-execution
24//!   helpers are available as methods on `GraphCompiler`, concrete input
25//!   slices/arrays, eager input slices/arrays, and tensor receivers.
26//!
27//! # Cargo features
28//!
29//! | Feature | Enables |
30//! |---|---|
31//! | `native` (default) and the other CPU provider features | Forwarded to `tenferro-cpu`; see its documentation. |
32//! | `autodiff` | The eager surface (`EagerSessionEinsumExt`) and AD rules. Adds the `tenferro-ad` dependency. |
33//! | `cuda` | CUDA execution through `tenferro-gpu`. |
34//! | `rocm` | Placeholder; HIP/ROCm is not implemented. |
35//!
36//! For contractions without AD, use the concrete [`TensorEinsumExt`] /
37//! [`TypedTensorEinsumExt`] APIs or [`ConcreteEinsumPlan`] inside a backend
38//! session; they need no `autodiff`.
39//!
40//! # Examples
41//!
42//! ```
43//! use tenferro_einsum::{ContractionTree, Subscripts};
44//!
45//! let subs = Subscripts::parse("ij,jk->ik").unwrap();
46//! let tree = ContractionTree::optimize(&subs, &[&[2, 3], &[3, 4]]).unwrap();
47//! assert_eq!(tree.step_count(), 1);
48//! ```
49//!
50//! ```
51//! use tenferro_cpu::CpuEngine;
52//! use tenferro_einsum::TensorEinsumExt;
53//! use tenferro_tensor::Tensor;
54//!
55//! let a = Tensor::from_vec_col_major(vec![2, 3], vec![1.0_f64; 6]).unwrap();
56//! let b = Tensor::from_vec_col_major(vec![3, 4], vec![1.0_f64; 12]).unwrap();
57//! let mut backend = CpuEngine::serial();
58//!
59//! let out = { let session = &mut backend;  {
60//!     [&a, &b].einsum("ij,jk->ik", session)
61//! } }?;
62//! assert_eq!(out.shape(), &[2, 4]);
63//! # Ok::<(), tenferro_einsum::Error>(())
64//! ```
65//!
66//! ```
67//! use tenferro_einsum::Subscripts;
68//!
69//! let trace = Subscripts::parse("ii->").unwrap();
70//! let diagonal = Subscripts::parse("ii->i").unwrap();
71//! let embedded = Subscripts::parse("i->ii").unwrap();
72//! let higher_rank = Subscripts::parse("iij->ij").unwrap();
73//!
74//! assert!(trace.output.is_empty());
75//! assert_eq!(diagonal.output, vec![b'i' as u32]);
76//! assert_eq!(embedded.output, vec![b'i' as u32, b'i' as u32]);
77//! assert_eq!(higher_rank.inputs[0], vec![b'i' as u32, b'i' as u32, b'j' as u32]);
78//! ```
79#![cfg_attr(docsrs, feature(doc_cfg))]
80// Without a CPU backend the concrete einsum route is unreachable by design and
81// only the typed `Unsupported` stub is live; the implementation it would call
82// is dead code in that configuration only.
83#![cfg_attr(
84    all(not(feature = "native"), not(feature = "blas")),
85    allow(dead_code, unused_imports)
86)]
87
88mod binary_dot;
89mod builder;
90mod cache;
91mod concrete;
92mod cpu_concrete;
93mod eager;
94#[cfg(feature = "autodiff")]
95mod eager_ad;
96mod ellipsis;
97mod error;
98mod extension;
99pub mod lowering;
100mod optimize;
101mod planning;
102pub mod prelude;
103mod subscripts;
104mod syntax;
105mod tensordot;
106mod traced;
107#[cfg(test)]
108mod typed_eager;
109pub(crate) mod util;
110
111pub use cache::EINSUM_EXTENSION_FAMILY_ID;
112pub use concrete::{
113    ConcreteEinsumPlan, TensorEinsumExt, TensorEinsumIntoExt, TensorReadEinsumExt,
114    TensorReadEinsumIntoExt, TensorTensordotExt, TypedTensorEinsumExt, TypedTensorEinsumIntoExt,
115    TypedTensorReadEinsumExt, TypedTensorReadEinsumIntoExt, TypedTensorTensordotExt,
116};
117#[cfg(feature = "autodiff")]
118#[cfg_attr(docsrs, doc(cfg(feature = "autodiff")))]
119pub use eager_ad::EagerSessionEinsumExt;
120pub use error::{Error, PlanningError, Result};
121pub use extension::extension_module;
122#[cfg(feature = "autodiff")]
123pub use extension::semantic_ad_rules;
124pub use optimize::EinsumOptimize;
125pub use planning::tree::{ContractionOptimizerOptions, ContractionTree};
126pub use subscripts::{
127    parse_einsum_notation, parse_einsum_subscripts, EinsumAxis, EinsumNotation, EinsumSubscripts,
128};
129pub use syntax::nested::NestedEinsum;
130pub use syntax::subscripts::Subscripts;
131pub use tensordot::TensorDotAxes;
132pub use traced::{TraceContextEinsumExt, TracedTensorEinsumExt};
133
134#[cfg(test)]
135mod concrete_tests;
136#[cfg(test)]
137mod tests;
138#[cfg(test)]
139mod typed_eager_tests;
tenferro_einsum/lib.rs

tenferro_einsum/
lib.rs
