Struct Einsum 

pub struct Einsum<L: Label = usize> {
    pub ixs: Vec<Vec<L>>,
    pub iy: Vec<L>,
    pub size_dict: HashMap<L, usize>,
    /* private fields */
}

Einsum specification and execution engine.

Supports contraction order optimization via omeco.

Example

use omeinsum::{Einsum, Tensor, Cpu};
use omeinsum::algebra::MaxPlus;
use std::collections::HashMap;

// A[i,j] × B[j,k] → C[i,k]
let a = Tensor::<f32, Cpu>::from_data(&[1.0, 2.0, 3.0, 4.0], &[2, 2]);
let b = Tensor::<f32, Cpu>::from_data(&[5.0, 6.0, 7.0, 8.0], &[2, 2]);

let sizes: HashMap<usize, usize> = [(0, 2), (1, 2), (2, 2)].into();
let mut ein = Einsum::new(
    vec![vec![0, 1], vec![1, 2]],
    vec![0, 2],
    sizes,
);

ein.optimize_greedy();
let result = ein.execute::<MaxPlus<f32>, f32, Cpu>(&[&a, &b]);
assert_eq!(result.shape(), &[2, 2]);

Fields

ixs: Vec<Vec<L>>

Input index labels for each tensor

iy: Vec<L>

Output index labels

size_dict: HashMap<L, usize>

Dimension sizes for each index

Implementations

impl<L: Label> Einsum<L>

pub fn new(ixs: Vec<Vec<L>>, iy: Vec<L>, size_dict: HashMap<L, usize>) -> Self

Create a new einsum specification.

Arguments

• ixs - Index labels for each input tensor
• iy - Output index labels
• size_dict - Mapping from index labels to dimension sizes

pub fn code(&self) -> EinCode<L>

Get the einsum code specification.

pub fn optimize_greedy(&mut self) -> &mut Self

Optimize contraction order using greedy algorithm.

Fast O(n²) algorithm, good for most cases.

pub fn optimize_treesa(&mut self) -> &mut Self

Optimize contraction order using simulated annealing.

Slower but finds better orderings for complex networks.

pub fn is_optimized(&self) -> bool

Check if optimization has been performed.

pub fn contraction_tree(&self) -> Option<&NestedEinsum<L>>

Get the optimized contraction tree.

impl Einsum<usize>

pub fn execute<A, T, B>(&self, tensors: &[&Tensor<T, B>]) -> Tensor<T, B>

where A: Algebra<Scalar = T, Index = u32>, T: Scalar + BackendScalar<B>, B: Backend + Default,

Execute the einsum contraction.

Type Parameters

• A - The algebra to use (e.g., Standard<f32>, MaxPlus<f32>)
• T - The scalar type
• B - The backend type

pub fn execute_with_argmax<A, T, B>( &self, tensors: &[&Tensor<T, B>], ) -> (Tensor<T, B>, Vec<Tensor<u32, B>>)

where A: Algebra<Scalar = T, Index = u32>, T: Scalar + BackendScalar<B>, B: Backend + Default,

Execute with argmax tracking for backpropagation.

Returns (result, argmax_cache) where argmax_cache contains argmax tensors for each binary contraction in the execution tree.