tenferro-einsum 0.9.0

Subscripts, contraction planning, concrete/traced/eager einsum APIs, extension runtime, and AD rule for tenferro.
Documentation 
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
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139

//! High-level einsum with N-ary contraction tree optimization.
//!
//! This crate provides:
//!
//! - **String notation**: `"ij,jk->ik"` (NumPy/PyTorch compatible)
//! - **Ellipsis notation**: `"...ij,...jk->...ik"` with right-aligned
//!   equal-or-one broadcasting, plus programmatic [`EinsumNotation`]
//! - **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.
//!
//! # Cargo features
//!
//! | Feature | Enables |
//! |---|---|
//! | `native` (default) and the other CPU provider features | Forwarded to `tenferro-cpu`; see its documentation. |
//! | `autodiff` | The eager surface (`EagerSessionEinsumExt`) and AD rules. Adds the `tenferro-ad` dependency. |
//! | `cuda` | CUDA execution through `tenferro-gpu`. |
//! | `rocm` | Placeholder; HIP/ROCm is not implemented. |
//!
//! For contractions without AD, use the concrete [`TensorEinsumExt`] /
//! [`TypedTensorEinsumExt`] APIs or [`ConcreteEinsumPlan`] inside a backend
//! session; they need no `autodiff`.
//!
//! # 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::CpuEngine;
//! 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 = CpuEngine::serial();
//!
//! let out = { let session = &mut backend;  {
//!     [&a, &b].einsum("ij,jk->ik", session)
//! } }?;
//! 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]);
//! ```
#![cfg_attr(docsrs, feature(doc_cfg))]
// Without a CPU backend the concrete einsum route is unreachable by design and
// only the typed `Unsupported` stub is live; the implementation it would call
// is dead code in that configuration only.
#![cfg_attr(
    all(not(feature = "native"), not(feature = "blas")),
    allow(dead_code, unused_imports)
)]

mod binary_dot;
mod builder;
mod cache;
mod concrete;
mod cpu_concrete;
mod eager;
#[cfg(feature = "autodiff")]
mod eager_ad;
mod ellipsis;
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")]
#[cfg_attr(docsrs, doc(cfg(feature = "autodiff")))]
pub use eager_ad::EagerSessionEinsumExt;
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_notation, parse_einsum_subscripts, EinsumAxis, EinsumNotation, 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;