Description
Julia's OMEinsum supports scaling factors in einsum operations via einsum!:
einsum!(rule, ix, iy, x, y, α, β) # y = α * einsum(x) + β * y
This allows in-place accumulation with scaling, which is important for:
- BLAS-style operations (GEMM uses α and β)
- Gradient accumulation in training loops
- Memory-efficient iterative algorithms
Current Behavior
Rust omeinsum-rs only supports the basic form:
let result = einsum::<Standard<f64>, _, _>(&[&a, &b], &[&[0, 1], &[1, 2]], &[0, 2]);
Proposed API
// Option 1: Separate function
einsum_scaled::<Standard<f64>, _, _>(
&[&a, &b],
&[&[0, 1], &[1, 2]],
&[0, 2],
alpha, // scale factor for result
beta, // scale factor for existing output
&mut output // in-place output
);
// Option 2: Builder pattern
ein("ij,jk->ik")
.alpha(2.0)
.beta(0.5)
.execute_into(&[&a, &b], &mut output);
Use Cases
- Gradient accumulation:
grad += learning_rate * new_grad
- BLAS compatibility: Match cuBLAS/MKL GEMM interface
- Iterative solvers:
x_new = α * A @ x + β * b
Implementation Notes
- cuTENSOR already supports alpha/beta (see
src/backend/cuda/cutensor/)
- CPU backend would need to add scaling in contraction loops
- Should work with all algebra types (Standard, Tropical)
Related
- Julia:
OMEinsum.jl/src/einsum.jl - unary_einsum!, binary_einsum!
- cuTENSOR: Already has
alpha/beta in contract API
Description
Julia's OMEinsum supports scaling factors in einsum operations via
einsum!:This allows in-place accumulation with scaling, which is important for:
Current Behavior
Rust omeinsum-rs only supports the basic form:
Proposed API
Use Cases
grad += learning_rate * new_gradx_new = α * A @ x + β * bImplementation Notes
src/backend/cuda/cutensor/)Related
OMEinsum.jl/src/einsum.jl-unary_einsum!,binary_einsum!alpha/betain contract API