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

//! High-level einsum with N-ary contraction tree optimization.
//!
//! This crate provides:
//!
//! - **String notation**: `"ij,jk->ik"` (NumPy/PyTorch compatible)
//! - **Parenthesized notation**: `"ij,(jk,kl)->il"` respects user-specified
//!   contraction order via [`NestedEinsum`]
//! - **Integer label notation**: using `u32` labels
//! - **Repeated labels**: `"ii->i"` extracts diagonals, `"ii->"` traces, and
//!   `"i->ii"` embeds a vector on a diagonal
//! - **N-ary contraction**: Automatic or manual optimization of pairwise
//!   contraction order via [`ContractionTree`]
//! - **Tensordot sugar**: NumPy-style axis-pair contraction extension methods,
//!   implemented as contraction sugar rather than as linear algebra APIs.
//! - **Concrete execution**: backend-explicit [`TensorEinsumExt`],
//!   [`TypedTensorEinsumExt`], [`TensorReadEinsumExt`],
//!   [`TypedTensorReadEinsumExt`], and [`ConcreteEinsumPlan`] APIs for non-AD
//!   tensor values.
//! - **Extension runtime**: traced einsum lowers to a registered tenferro
//!   extension runtime, keeping core op definitions small.
//! - **Tensor extension traits**: graph-building and immediate-execution
//!   helpers are available as methods on `GraphCompiler`, concrete input
//!   slices/arrays, eager input slices/arrays, and tensor receivers.
//!
//! # Examples
//!
//! ```
//! use tenferro_einsum::{ContractionTree, Subscripts};
//!
//! let subs = Subscripts::parse("ij,jk->ik").unwrap();
//! let tree = ContractionTree::optimize(&subs, &[&[2, 3], &[3, 4]]).unwrap();
//! assert_eq!(tree.step_count(), 1);
//! ```
//!
//! ```
//! use tenferro_cpu::CpuBackend;
//! use tenferro_einsum::TensorEinsumExt;
//! use tenferro_tensor::Tensor;
//!
//! let a = Tensor::from_vec_col_major(vec![2, 3], vec![1.0_f64; 6]).unwrap();
//! let b = Tensor::from_vec_col_major(vec![3, 4], vec![1.0_f64; 12]).unwrap();
//! let mut backend = CpuBackend::new();
//!
//! let out = [&a, &b].einsum("ij,jk->ik", &mut backend)?;
//! assert_eq!(out.shape(), &[2, 4]);
//! # Ok::<(), tenferro_einsum::Error>(())
//! ```
//!
//! ```
//! use tenferro_einsum::Subscripts;
//!
//! let trace = Subscripts::parse("ii->").unwrap();
//! let diagonal = Subscripts::parse("ii->i").unwrap();
//! let embedded = Subscripts::parse("i->ii").unwrap();
//! let higher_rank = Subscripts::parse("iij->ij").unwrap();
//!
//! assert!(trace.output.is_empty());
//! assert_eq!(diagonal.output, vec![b'i' as u32]);
//! assert_eq!(embedded.output, vec![b'i' as u32, b'i' as u32]);
//! assert_eq!(higher_rank.inputs[0], vec![b'i' as u32, b'i' as u32, b'j' as u32]);
//! ```

mod binary_dot;
mod builder;
mod cache;
mod concrete;
mod eager;
#[cfg(feature = "autodiff")]
mod eager_ad;
mod error;
mod extension;
pub mod lowering;
mod optimize;
mod planning;
pub mod prelude;
mod subscripts;
mod syntax;
mod tensordot;
mod traced;
#[cfg(test)]
mod typed_eager;
pub(crate) mod util;

pub use cache::EINSUM_EXTENSION_FAMILY_ID;
pub use concrete::{
    ConcreteEinsumPlan, TensorEinsumExt, TensorEinsumIntoExt, TensorReadEinsumExt,
    TensorReadEinsumIntoExt, TensorTensordotExt, TypedTensorEinsumExt, TypedTensorEinsumIntoExt,
    TypedTensorReadEinsumExt, TypedTensorReadEinsumIntoExt, TypedTensorTensordotExt,
};
#[cfg(feature = "autodiff")]
pub use eager_ad::{EagerEinsumExt, EagerTensorEinsumExt};
pub use error::{Error, PlanningError, Result};
pub use extension::extension_module;
#[cfg(feature = "autodiff")]
pub use extension::semantic_ad_rules;
pub use optimize::EinsumOptimize;
pub use planning::tree::{ContractionOptimizerOptions, ContractionTree};
pub use subscripts::{parse_einsum_subscripts, EinsumSubscripts};
pub use syntax::nested::NestedEinsum;
pub use syntax::subscripts::Subscripts;
pub use tensordot::TensorDotAxes;
pub use traced::{TraceContextEinsumExt, TracedTensorEinsumExt};

#[cfg(test)]
mod concrete_tests;
#[cfg(test)]
mod tests;
#[cfg(test)]
mod typed_eager_tests;