Skip to main content

Module eager_einsum

Module eager_einsum 

Source
Expand description

Eager einsum helpers exposed through the tenferro facade.

This module provides both immediate eager execution on concrete tensors and eager reverse-mode autodiff on tracked tensors.

§Examples

use tenferro::eager_einsum::{eager_einsum, eager_einsum_ad};
use tenferro::{CpuBackend, EagerTensor, Tensor};

let mut backend = CpuBackend::new();
let a = Tensor::from_vec(vec![3], vec![1.0_f64, 2.0, 3.0]);
let b = Tensor::from_vec(vec![3], vec![4.0_f64, 5.0, 6.0]);
let dot = eager_einsum(&mut backend, &[&a, &b], "i,i->").unwrap();

assert_eq!(dot.as_slice::<f64>().unwrap(), &[32.0]);

let x = EagerTensor::requires_grad(Tensor::from_vec(vec![3], vec![1.0_f64, 2.0, 3.0]));
let y = EagerTensor::requires_grad(Tensor::from_vec(vec![3], vec![4.0_f64, 5.0, 6.0]));
let loss = eager_einsum_ad(&[&x, &y], "i,i->").unwrap();
let _ = loss.backward().unwrap();

assert_eq!(x.grad().unwrap().as_slice::<f64>().unwrap(), &[4.0, 5.0, 6.0]);
assert_eq!(y.grad().unwrap().as_slice::<f64>().unwrap(), &[1.0, 2.0, 3.0]);

Functions§

eager_einsum
Eager N-ary einsum on concrete Tensor values.
eager_einsum_ad
Execute an einsum eagerly and record it for reverse-mode autodiff.