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

• einsum-derive crate is proc-macro crate to provide above einsum! macro.
• einsum-codegen crate implements parser for the einsum subscripts like ij,jk->ik and generates Rust code.

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

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.