prepare_contraction

Function prepare_contraction 

Source
pub fn prepare_contraction<I: IndexLike>(
    indices_a: &[I],
    dims_a: &[usize],
    indices_b: &[I],
    dims_b: &[usize],
) -> Result<ContractionSpec<I>, ContractionError>
Expand description

Prepare contraction data for two tensors that share common indices.

This function finds common indices and computes the axes to contract and the resulting indices/dimensions.

ยงExample

use tensor4all_core::index::DefaultIndex as Index;
use tensor4all_core::index_ops::prepare_contraction;

let i = Index::new_dyn(2);
let j = Index::new_dyn(3);
let k = Index::new_dyn(4);

let indices_a = vec![i.clone(), j.clone()];
let dims_a = vec![2, 3];
let indices_b = vec![j.clone(), k.clone()];
let dims_b = vec![3, 4];

let spec = prepare_contraction(&indices_a, &dims_a, &indices_b, &dims_b).unwrap();
assert_eq!(spec.axes_a, vec![1]);  // j is at position 1 in a
assert_eq!(spec.axes_b, vec![0]);  // j is at position 0 in b
assert_eq!(spec.result_dims, vec![2, 4]);  // [i, k]