Struct einsum_codegen::Subscripts

source ·

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

Expand description

Einsum subscripts with tensor names, e.g. ij,jk->ik | arg0 arg1 -> out

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::*};

let mut names = Namespace::init();

// Infer output subscripts for implicit mode
let raw = RawSubscripts::from_str("ij,jk").unwrap();
let subscripts = Subscripts::from_raw(&mut names, raw);
assert_eq!(subscripts.output.raw(), &['i', 'k']);

// Reordered alphabetically
let raw = RawSubscripts::from_str("ji").unwrap();
let subscripts = Subscripts::from_raw(&mut names, raw);
assert_eq!(subscripts.output.raw(), &['i', 'j']);

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, "ij,jk->ik").unwrap();
assert_eq!(subscripts.contraction_indices(), btreeset!{'j'});

// Reduce all Tr(AB)
let subscripts = Subscripts::from_raw_indices(&mut names, "ij,ji->").unwrap();
assert_eq!(subscripts.contraction_indices(), btreeset!{'i', 'j'});

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

source

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

Factorize subscripts

ij,jk,kl->il | arg0 arg1 arg2 -> out0

will be factorized with (arg0, arg1) into

ij,jk->ik | arg0 arg1 -> out1
ik,kl->il | out1 arg2 -> out0

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, "ij,jk,kl->il").unwrap();

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