Skip to main content

slice_code

Function slice_code 

Source
pub fn slice_code<L: Label>(
    code: &NestedEinsum<L>,
    size_dict: &HashMap<L, usize>,
    config: &TreeSASlicer,
    original_ixs: &[Vec<L>],
) -> Option<SlicedEinsum<L>>
Expand description

Slice a NestedEinsum to reduce space complexity.

This is the main entry point for slicing optimization. It takes an already-optimized contraction tree and finds indices to slice over to reduce memory requirements.

§Arguments

  • code - The optimized NestedEinsum to slice
  • size_dict - Size dictionary mapping labels to their dimensions
  • config - Slicing configuration
  • original_ixs - Original input tensor indices (from the EinCode)

§Returns

A SlicedEinsum containing the sliced indices and the optimized tree.

§Example

use omeco::{EinCode, TreeSASlicer, optimize_code, slice_code, GreedyMethod, uniform_size_dict};

let code = EinCode::new(
    vec![vec!['i', 'j'], vec!['j', 'k']],
    vec!['i', 'k'],
);
let sizes = uniform_size_dict(&code, 64);

let optimized = optimize_code(&code, &sizes, &GreedyMethod::default()).unwrap();
let slicer = TreeSASlicer::fast().with_sc_target(5.0);
let sliced = slice_code(&optimized, &sizes, &slicer, &code.ixs);