Skip to main content

Crate tenferro_einsum

Crate tenferro_einsum 

Source
Expand description

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.

§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::{BackendSessionHost, 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 = backend.with_backend_session(|session| {
    [&a, &b].einsum("ij,jk->ik", session)
})?;
assert_eq!(out.shape(), &[2, 4]);
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]);

Modules§

lowering
Read-only einsum lowering plans.
prelude
Einsum and tensordot extension traits for concrete, eager, and traced values.

Structs§

ConcreteEinsumPlan
Prepared concrete einsum plan for repeated executions with fixed input dtype and shape metadata.
ContractionOptimizerOptions
Public options for automatic contraction-path optimization.
ContractionTree
Contraction tree determining pairwise contraction order for N-ary einsum.
EinsumNotation
Rank-unresolved einsum notation.
EinsumSubscripts
Canonical N-ary einsum subscripts using integer labels.
Subscripts
Einsum subscripts using integer labels (omeinsum-rs compatible).

Enums§

EinsumAxis
One unresolved axis token in rank-polymorphic einsum notation.
EinsumOptimize
Controls how the contraction path is determined for N-ary einsum.
Error
Errors produced while parsing, planning, lowering, or executing einsum expressions.
NestedEinsum
Recursive einsum tree that preserves parenthesized grouping.
PlanningError
Domain-specific cause of an einsum planning failure.
TensorDotAxes
Axis specification for TracedTensorEinsumExt::tensordot contraction sugar.

Constants§

EINSUM_EXTENSION_FAMILY_ID
Stable family identifier for the standard tenferro einsum extension.

Traits§

EagerEinsumExt
Eager einsum extension methods for slices or arrays of EagerTensor refs.
EagerTensorEinsumExt
Eager tensor contraction-sugar methods.
TensorEinsumExt
Backend-explicit einsum methods for dtype-erased concrete tensors.
TensorEinsumIntoExt
Backend-explicit preallocated-output einsum methods for dtype-erased tensors.
TensorReadEinsumExt
Backend-explicit einsum methods for TensorRead inputs.
TensorReadEinsumIntoExt
Backend-explicit preallocated-output einsum methods for TensorRead inputs.
TensorTensordotExt
Backend-explicit tensordot sugar for dtype-erased concrete tensors.
TraceContextEinsumExt
Backend-neutral einsum tracing methods for TraceContext.
TracedTensorEinsumExt
Traced tensor contraction-sugar methods.
TypedTensorEinsumExt
Backend-explicit einsum methods for typed concrete tensors.
TypedTensorEinsumIntoExt
Backend-explicit preallocated-output einsum methods for typed concrete tensors.
TypedTensorReadEinsumExt
Backend-explicit einsum methods for typed borrowed views.
TypedTensorReadEinsumIntoExt
Backend-explicit preallocated-output einsum methods for typed borrowed views.
TypedTensorTensordotExt
Backend-explicit tensordot sugar for typed concrete tensors.

Functions§

extension_module
Build this extension module for one runtime engine.
parse_einsum_notation
Parse string notation into rank-unresolved axis tokens.
parse_einsum_subscripts
Parse string einsum notation into canonical integer labels.
semantic_ad_rules
Return the semantic-program einsum extension AD rules.

Type Aliases§

Result
Result type alias for einsum parsing, planning, and all public extension APIs.