Struct einsum_codegen::Subscripts

source ·

pub struct Subscripts {
    pub inputs: Vec<Subscript>,
    pub output: Subscript,
}

Expand description

Einsum subscripts with tensor names, e.g. ab,bc->ac | arg0,arg1->out0

Indices are remapped as starting from a to distinguish same subscripts, e.g. i,i-> and j,j->

use einsum_codegen::{*, parser::RawSubscript};

let mut names = Namespace::init();
let mut ss1 = Subscripts::from_raw_indices(&mut names, "ij,jk,kl->il").unwrap();

let mut names = Namespace::init();
let mut ss2 = Subscripts::from_raw_indices(&mut names, "xz,zy,yw->xw").unwrap();

assert_eq!(ss1, ss2);
assert_eq!(ss1.to_string(), "ab,bc,cd->ad | arg0,arg1,arg2->out0");
assert_eq!(ss2.to_string(), "ab,bc,cd->ad | arg0,arg1,arg2->out0");

Fields§

§inputs: Vec<Subscript>

Input subscript, ij and jk

§output: Subscript

Output subscript.

Implementations§

source §

impl Subscripts

source

pub fn compute_order(&self) -> usize

Returns $\alpha$ if this subscripts requires $O(N^\alpha)$ floating point operation

source

pub fn memory_order(&self) -> usize

Returns $\beta$ if this subscripts requires $O(N^\beta)$ memory

source

pub fn from_raw(names: &mut Namespace, raw: RawSubscripts) -> Self

Normalize subscripts into “explicit mode”

numpy.einsum has “explicit mode” including ->, e.g. ij,jk->ik and “implicit mode” e.g. ij,jk. The output subscript is determined from input subscripts in implicit mode:

In implicit mode, the chosen subscripts are important since the axes of the output are reordered alphabetically. This means that np.einsum('ij', a) doesn’t affect a 2D array, while np.einsum('ji', a) takes its transpose. Additionally, np.einsum('ij,jk', a, b) returns a matrix multiplication, while, np.einsum('ij,jh', a, b) returns the transpose of the multiplication since subscript ‘h’ precedes subscript ‘i’.

use std::str::FromStr;
use einsum_codegen::{*, parser::*};

// Infer output subscripts for implicit mode
let mut names = Namespace::init();
let raw = RawSubscripts::from_str("ab,bc").unwrap();
let subscripts = Subscripts::from_raw(&mut names, raw);
assert_eq!(subscripts.to_string(), "ab,bc->ac | arg0,arg1->out0");

// Reordered alphabetically
let mut names = Namespace::init(); // reset namespace
let raw = RawSubscripts::from_str("ba").unwrap();
let subscripts = Subscripts::from_raw(&mut names, raw);
assert_eq!(subscripts.to_string(), "ab->ba | arg0->out0");

source

pub fn from_raw_indices(names: &mut Namespace, indices: &str) -> Result<Self>

source

pub fn contraction_indices(&self) -> BTreeSet<char>

Indices to be contracted

use std::str::FromStr;
use maplit::btreeset;
use einsum_codegen::*;

let mut names = Namespace::init();

// Matrix multiplication AB
let subscripts = Subscripts::from_raw_indices(&mut names, "ab,bc->ac").unwrap();
assert_eq!(subscripts.contraction_indices(), btreeset!{'b'});

// Reduce all Tr(AB)
let subscripts = Subscripts::from_raw_indices(&mut names, "ab,ba->").unwrap();
assert_eq!(subscripts.contraction_indices(), btreeset!{'a', 'b'});

// Take diagonal elements
let subscripts = Subscripts::from_raw_indices(&mut names, "aa->a").unwrap();
assert_eq!(subscripts.contraction_indices(), btreeset!{});

source

pub fn factorize( &self, names: &mut Namespace, inners: BTreeSet<Position> ) -> Result<(Self, Self)>

Factorize subscripts

ab,bc,cd->ad | arg0,arg1,arg2->out0

will be factorized with (arg0, arg1) into

ab,bc->ac | arg0,arg1 -> out1
ab,bc->ac | out1 arg2 -> out0

Be sure that the indices of out1 in the first step ac is renamed into ab in the second step.

use einsum_codegen::{*, parser::RawSubscript};
use std::str::FromStr;
use maplit::btreeset;

let mut names = Namespace::init();
let base = Subscripts::from_raw_indices(&mut names, "ab,bc,cd->ad").unwrap();

let (step1, step2) = base.factorize(&mut names,
  btreeset!{ Position::Arg(0), Position::Arg(1) }
).unwrap();

assert_eq!(step1.to_string(), "ab,bc->ac | arg0,arg1->out1");
assert_eq!(step2.to_string(), "ab,bc->ac | out1,arg2->out0");