Skip to main content

candle_einops_macros/
lib.rs

1mod einops;
2mod einsum;
3
4/// Macro to perform tensor transformations using simple expressions
5///
6/// This macro is re-exported as `candle_einops::einops`; the runtime crate's
7/// documentation contains a complete runnable example.
8#[proc_macro]
9pub fn einops(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
10    einops::einops(input.into())
11        .unwrap_or_else(|e| e.to_compile_error())
12        .into()
13}
14
15/// Evaluates an explicit-output Einstein summation equation.
16///
17/// Supports any positive number of operands with named axes and at most one
18/// `..` per axis list. Repeated input labels extract diagonals before
19/// contraction. Use this through `candle_einops::einsum`; generated code calls
20/// a private runtime ABI, so `candle-einops-macros` and `candle-einops` must use
21/// exactly the same version.
22#[proc_macro]
23pub fn einsum(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
24    einsum::einsum(input.into())
25        .unwrap_or_else(|error| error.to_compile_error())
26        .into()
27}