omeinsum/
lib.rs

1//! # OMEinsum-rs
2//!
3//! High-performance Einstein summation with support for both tropical and standard algebras.
4//!
5//! ## Features
6//!
7//! - **Algebra-agnostic**: Works with standard arithmetic `(+, ×)` and tropical semirings `(max, +)`, `(min, +)`
8//! - **Optimized contraction**: Integration with [omeco](https://github.com/GiggleLiu/omeco) for contraction order optimization
9//! - **Backpropagation**: Gradient computation for both tropical and standard operations
10//! - **Zero-copy views**: Stride-based tensor with efficient permute/reshape
11//! - **CPU + CUDA**: Support for both backends (CUDA optional)
12//!
13//! ## Quick Start
14//!
15//! ```rust
16//! use omeinsum::{Tensor, Einsum, einsum, Cpu};
17//! use omeinsum::algebra::{Standard, MaxPlus};
18//!
19//! // Standard matrix multiplication
20//! let a = Tensor::<f32, Cpu>::from_data(&[1.0, 2.0, 3.0, 4.0], &[2, 2]);
21//! let b = Tensor::<f32, Cpu>::from_data(&[1.0, 2.0, 3.0, 4.0], &[2, 2]);
22//!
23//! // C[i,k] = Σ_j A[i,j] × B[j,k]
24//! let c = einsum::<Standard<f32>, _, _>(&[&a, &b], &[&[0, 1], &[1, 2]], &[0, 2]);
25//!
26//! // Tropical (max-plus) matrix multiplication
27//! // C[i,k] = max_j (A[i,j] + B[j,k])
28//! let c_tropical = einsum::<MaxPlus<f32>, _, _>(&[&a, &b], &[&[0, 1], &[1, 2]], &[0, 2]);
29//! ```
30//!
31//! ## Architecture
32//!
33//! ```text
34//! ┌─────────────────────────────────────────────────────────────┐
35//! │                         User API                            │
36//! │   einsum(tensors, ixs, iy) → Tensor                        │
37//! │   Einsum::new(ixs, iy).optimize().execute(tensors)         │
38//! └─────────────────────────────────────────────────────────────┘
39//!                               │
40//!                               ▼
41//! ┌─────────────────────────────────────────────────────────────┐
42//! │                      Einsum Engine                          │
43//! │   omeco::optimize_code() → NestedEinsum (contraction tree) │
44//! │   Execute tree via binary contractions                      │
45//! └─────────────────────────────────────────────────────────────┘
46//!                               │
47//!                               ▼
48//! ┌─────────────────────────────────────────────────────────────┐
49//! │                 Algebra<T> dispatch                         │
50//! │   Standard<T>: (+, ×) → BLAS/loops                         │
51//! │   MaxPlus<T>:  (max, +) → tropical-gemm                    │
52//! │   MinPlus<T>:  (min, +) → tropical-gemm                    │
53//! └─────────────────────────────────────────────────────────────┘
54//! ```
55
56pub mod algebra;
57pub mod backend;
58pub mod einsum;
59pub mod tensor;
60
61// Re-exports
62pub use algebra::{Algebra, Complex32, Complex64, Semiring, Standard};
63pub use backend::{Backend, BackendScalar, Cpu, Storage};
64pub use einsum::{einsum, einsum_with_grad, EinBuilder, Einsum};
65pub use tensor::{Tensor, TensorView};
66
67#[cfg(feature = "tropical")]
68pub use algebra::{MaxMul, MaxPlus, MinPlus};
69
70#[cfg(feature = "cuda")]
71pub use backend::Cuda;
omeinsum/lib.rs

omeinsum/
lib.rs
