termoshtt/einsum-derive

Proc-macro based einsum implementation for rust-ndarray

einsum-derive

Proc-macro based einsum implementation for ndarray crate

use ndarray::array;
use einsum_derive::einsum;

let a = array![
  [1.0, 2.0],
  [3.0, 4.0]
];
let b = array![
  [1.0, 2.0],
  [3.0, 4.0]
];
let c = einsum!("ij,jk->ik", a, b);
assert_eq!(c, array![
  [6.0, 8.0],
  [12.0, 16.0]
]);

This proc-macro wil compile the input subscripts "ij,jk->ik" to generate Rust code executing corresponding operation.

Status / Roadmap

• Optimal contraction by memorizing partial summation to reduce computation order.
  • For example, three matrix multiplication ij,jk,kl->il is factorized into two successive einsum ij,jk->ik and ik,kl->il.
• Call BLAS routines if possible
• Ellipsis ... support

Architecture

	crates.io	docs.rs	GitHub Pages	Description
einsum-derive				proc-macro crate to provide einsum! macro
einsum-codegen				Implements parser for the einsum subscripts and generates Rust code

Benchmark

Benchmark with criterion.rs is running on GitHub Action on every commit on the main branch. The code is placed at einsum-derive/benches/einsum.rs, and you can run it on your environment by

cargo bench

and you will find its result on target/criterion/report/index.html.

License

© 2022 Toshiki Teramura (@termoshtt)

This project is licensed under either of

• Apache License, Version 2.0, (LICENSE-APACHE or https://www.apache.org/licenses/LICENSE-2.0)
• MIT license (LICENSE-MIT or https://opensource.org/licenses/MIT)

at your option.

Links

• numpy.einsum is well-known einsum implementation in Python.
• opt_einsum is an implementation for optimizing einsum computation for NumPy and other linear algebra packages.
• oracleofnj/einsum is a runtime-based implementation of einsum for rust-ndarray