Struct ContractionTree 

pub struct ContractionTree { /* private fields */ }

Contraction tree determining pairwise contraction order for N-ary einsum.

When contracting more than two tensors, the order in which pairwise contractions are performed significantly affects performance. ContractionTree encodes this order as a binary tree.

Optimization

Use ContractionTree::optimize for automatic cost-based optimization (e.g., greedy algorithm based on tensor sizes), or ContractionTree::from_pairs for manual specification.

Implementations

impl ContractionTree

pub fn optimize(subscripts: &Subscripts, shapes: &[&[usize]]) -> Result<Self>

Automatically compute an optimized contraction order.

Uses a cost-based heuristic (greedy algorithm) to determine the pairwise contraction sequence that minimizes total operation count.

Arguments

• subscripts — Einsum subscripts for all tensors
• shapes — Shape of each input tensor

Errors

Returns an error if subscripts and shapes are inconsistent.

pub fn optimize_with_options( subscripts: &Subscripts, shapes: &[&[usize]], options: &ContractionOptimizerOptions, ) -> Result<Self>

Automatically compute an optimized contraction order with explicit planner options.

This routes automatic planning through TreeSA using the provided configuration. The default options correspond to a greedy-initialized TreeSA with zero annealing iterations.

Errors

Returns an error if subscripts, shapes, or planner options are invalid.

pub fn from_pairs( subscripts: &Subscripts, shapes: &[&[usize]], pairs: &[(usize, usize)], ) -> Result<Self>

Manually build a contraction tree from a pairwise contraction sequence.

Each pair (i, j) specifies which two tensors (or intermediate results) to contract next. Intermediate results are assigned indices starting from the number of input tensors.

Arguments

• subscripts — Einsum subscripts for all tensors
• shapes — Shape of each input tensor
• pairs — Ordered list of pairwise contractions

Examples

// Three tensors: A[ij] B[jk] C[kl] -> D[il]
// Contract B and C first, then A with the result:
let subs = Subscripts::new(&[&[0, 1], &[1, 2], &[2, 3]], &[0, 3]);
let shapes = [&[3, 4][..], &[4, 5], &[5, 6]];
let tree = ContractionTree::from_pairs(
    &subs,
    &shapes,
    &[(1, 2), (0, 3)],  // B*C -> T(index=3), then A*T -> D
).unwrap();

Errors

Returns an error if the pairs do not form a valid contraction sequence.

pub fn step_count(&self) -> usize

Return the number of pairwise contraction steps in this tree.

Examples

use tenferro_einsum::{ContractionTree, Subscripts};

let subs = Subscripts::new(&[&[0, 1], &[1, 2], &[2, 3]], &[0, 3]);
let tree = ContractionTree::from_pairs(
    &subs,
    &[&[2, 2], &[2, 2], &[2, 2]],
    &[(1, 2), (0, 3)],
)
.unwrap();
assert_eq!(tree.step_count(), 2);

pub fn step_pair(&self, step_idx: usize) -> Option<(usize, usize)>

Return the operand indices for a pairwise contraction step.

The returned indices refer to the original inputs (0..n_inputs) and then to intermediates (n_inputs..) produced by earlier steps.

Examples

use tenferro_einsum::{ContractionTree, Subscripts};

let subs = Subscripts::new(&[&[0, 1], &[1, 2], &[2, 3]], &[0, 3]);
let tree = ContractionTree::from_pairs(
    &subs,
    &[&[2, 2], &[2, 2], &[2, 2]],
    &[(1, 2), (0, 3)],
)
.unwrap();
assert_eq!(tree.step_pair(0), Some((1, 2)));

pub fn step_subscripts( &self, step_idx: usize, ) -> Option<(&[u32], &[u32], &[u32])>

Return the (lhs, rhs, output) subscripts for a pairwise step.

The output subscripts are the intermediate labels preserved after the contraction, or the final output labels on the last step.

Examples

use tenferro_einsum::{ContractionTree, Subscripts};

let subs = Subscripts::new(&[&[0, 1], &[1, 2], &[2, 3]], &[0, 3]);
let tree = ContractionTree::from_pairs(
    &subs,
    &[&[2, 2], &[2, 2], &[2, 2]],
    &[(1, 2), (0, 3)],
)
.unwrap();
let (lhs, rhs, out) = tree.step_subscripts(0).unwrap();
assert_eq!(lhs, &[1, 2]);
assert_eq!(rhs, &[2, 3]);
assert_eq!(out, &[1, 3]);