Skip to content

Add alpha/beta scaling factors for einsum operations #23

Description

@GiggleLiu

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

  1. Gradient accumulation: grad += learning_rate * new_grad
  2. BLAS compatibility: Match cuBLAS/MKL GEMM interface
  3. 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

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions