tenferro-einsum 0.8.0

use std::sync::Arc;

use tenferro_ad::{EagerRuntime, EagerTensor};
use tenferro_cpu::CpuEngine;
use tenferro_runtime::ExtensionCacheSelector;
use tenferro_tensor::{CacheStats, Tensor, TensorRead, TensorView};

use super::backend_broadcast_multiply_untracked;

/// Run one eager einsum inside a session on the caller's engine.
fn einsum_with_engine(
    engine: &mut CpuEngine,
    inputs: &[&EagerTensor],
    subscripts: &str,
) -> crate::Result<EagerTensor> {
    inputs[0]
        .runtime()
        .with_eager_session(engine, |s| super::einsum(s, inputs, subscripts))
}

/// One eager session per call on a fresh serial engine.
fn einsum(inputs: &[&EagerTensor], subscripts: &str) -> crate::Result<EagerTensor> {
    einsum_with_engine(&mut CpuEngine::serial(), inputs, subscripts)
}

fn extension_cache_stats(engine: &mut CpuEngine, ctx: &Arc<EagerRuntime>) -> CacheStats {
    ctx.with_eager_session(engine, |s| {
        s.with_extension_caches(|caches| caches.stats(ExtensionCacheSelector::All))
    })
    .unwrap()
}

#[test]
fn dot_general_retained_bytes_count_only_spilled_capacity() {
    use tenferro_tensor::DotGeneralConfig;
    for count in [0, 4, 5, 9] {
        let config = DotGeneralConfig {
            lhs_contracting_dims: (0..count).collect(),
            rhs_contracting_dims: (0..count).rev().collect(),
            lhs_batch_dims: (count..2 * count).collect(),
            rhs_batch_dims: (count..2 * count).rev().collect(),
        };
        let expected: usize = [
            &config.lhs_contracting_dims,
            &config.rhs_contracting_dims,
            &config.lhs_batch_dims,
            &config.rhs_batch_dims,
        ]
        .into_iter()
        .filter(|axes| axes.is_heap())
        .map(|axes| axes.capacity() * size_of::<usize>())
        .sum();
        assert_eq!(
            super::std_tensor_op_retained_bytes(&super::StdTensorOp::DotGeneral { config }),
            expected
        );
        if count <= 4 {
            assert_eq!(expected, 0);
        }
    }
}

#[test]
fn binary_einsum_col_major_matmul_uses_direct_dot_general_fast_path() {
    let ctx = EagerRuntime::new().unwrap();
    let lhs = EagerTensor::from_tensor_in(
        Tensor::from_vec_col_major(vec![2, 3], vec![1.0_f64; 6]).unwrap(),
        ctx.clone(),
    )
    .unwrap();
    let rhs = EagerTensor::from_tensor_in(
        Tensor::from_vec_col_major(vec![4, 2], vec![1.0_f64; 8]).unwrap(),
        ctx.clone(),
    )
    .unwrap();

    let mut engine = CpuEngine::serial();
    let out = einsum_with_engine(&mut engine, &[&lhs, &rhs], "ji,kj->ki").unwrap();

    assert_eq!(out.shape(), &[4, 3]);
    assert_eq!(
        out.value().unwrap().as_slice::<f64>().unwrap(),
        &[2.0_f64; 12]
    );
    assert_eq!(extension_cache_stats(&mut engine, &ctx).entries, 0);
}

#[test]
fn outer_product_einsum_reuses_its_extension_cache_entry() {
    let ctx = EagerRuntime::new().unwrap();
    let a = EagerTensor::from_tensor_in(
        Tensor::from_vec_col_major(vec![2, 3], (0..6).map(|i| i as f64 + 1.0).collect()).unwrap(),
        ctx.clone(),
    )
    .unwrap();
    let b = EagerTensor::from_tensor_in(
        Tensor::from_vec_col_major(vec![3, 4], (0..12).map(|i| i as f64 * 0.5 - 2.0).collect())
            .unwrap(),
        ctx.clone(),
    )
    .unwrap();
    let mut engine = CpuEngine::serial();
    let first = einsum_with_engine(&mut engine, &[&a, &b], "ij,kl->ijkl").unwrap();
    let first_stats = extension_cache_stats(&mut engine, &ctx);
    let second = einsum_with_engine(&mut engine, &[&a, &b], "ij,kl->ijkl").unwrap();
    let second_stats = extension_cache_stats(&mut engine, &ctx);
    assert_eq!(first.shape(), second.shape());
    assert!(first_stats.entries > 0);
    assert_eq!(first_stats.entries, second_stats.entries);
    assert_eq!(second_stats.hits, first_stats.hits + 1);
}

#[test]
fn nary_eager_einsum_expands_to_standard_ops_with_runtime_cache() {
    let ctx = EagerRuntime::new().unwrap();
    let a = EagerTensor::from_tensor_in(
        Tensor::from_vec_col_major(vec![2, 3], vec![1.0_f64; 6]).unwrap(),
        ctx.clone(),
    )
    .unwrap();
    let b = EagerTensor::from_tensor_in(
        Tensor::from_vec_col_major(vec![3, 4], vec![1.0_f64; 12]).unwrap(),
        ctx.clone(),
    )
    .unwrap();
    let c = EagerTensor::from_tensor_in(
        Tensor::from_vec_col_major(vec![4, 5], vec![1.0_f64; 20]).unwrap(),
        ctx.clone(),
    )
    .unwrap();

    let mut engine = CpuEngine::serial();
    let out = einsum_with_engine(&mut engine, &[&a, &b, &c], "ij,jk,kl->il").unwrap();

    assert_eq!(out.shape(), &[2, 5]);
    assert_eq!(
        out.value().unwrap().as_slice::<f64>().unwrap(),
        &[12.0_f64; 10]
    );
    assert_eq!(extension_cache_stats(&mut engine, &ctx).entries, 1);
}

#[test]
fn nary_eager_einsum_expanded_standard_ops_reuse_runtime_cache() {
    let ctx = EagerRuntime::new().unwrap();
    let a = EagerTensor::from_tensor_in(
        Tensor::from_vec_col_major(vec![2, 3], vec![1.0_f64; 6]).unwrap(),
        ctx.clone(),
    )
    .unwrap();
    let b = EagerTensor::from_tensor_in(
        Tensor::from_vec_col_major(vec![3, 4], vec![1.0_f64; 12]).unwrap(),
        ctx.clone(),
    )
    .unwrap();
    let c = EagerTensor::from_tensor_in(
        Tensor::from_vec_col_major(vec![4, 5], vec![1.0_f64; 20]).unwrap(),
        ctx.clone(),
    )
    .unwrap();

    let mut engine = CpuEngine::serial();
    let first = einsum_with_engine(&mut engine, &[&a, &b, &c], "ij,jk,kl->il").unwrap();

    assert_eq!(first.shape(), &[2, 5]);
    let after_first = extension_cache_stats(&mut engine, &ctx);
    assert_eq!(after_first.entries, 1);
    assert!(after_first.retained_bytes > 0);

    let second = einsum_with_engine(&mut engine, &[&a, &b, &c], "ij,jk,kl->il").unwrap();

    assert_eq!(second.shape(), &[2, 5]);
    let after_second = extension_cache_stats(&mut engine, &ctx);
    assert_eq!(after_second.entries, 1);
    assert_eq!(after_second.retained_bytes, after_first.retained_bytes);

    ctx.with_eager_session(&mut engine, |s| {
        s.with_extension_caches(|caches| caches.clear())
    })
    .unwrap();
    assert_eq!(extension_cache_stats(&mut engine, &ctx).entries, 0);
}

#[test]
fn expanded_eager_einsum_cache_hit_preserves_lazy_view_output() {
    let ctx = EagerRuntime::new().unwrap();
    let k = 2;
    let j = 3;
    let o = 4;
    let t = 5;
    let lhs_data: Vec<f64> = (0..k * j * t).map(|idx| idx as f64 + 1.0).collect();
    let rhs_data: Vec<f64> = (0..o * t).map(|idx| idx as f64 + 101.0).collect();
    let lhs = EagerTensor::from_tensor_in(
        Tensor::from_vec_col_major(vec![k, j, t], lhs_data).unwrap(),
        ctx.clone(),
    )
    .unwrap();
    let rhs = EagerTensor::from_tensor_in(
        Tensor::from_vec_col_major(vec![o, t], rhs_data).unwrap(),
        ctx.clone(),
    )
    .unwrap();

    let mut engine = CpuEngine::serial();
    let first = einsum_with_engine(&mut engine, &[&lhs, &rhs], "kjt,ot->jkot").unwrap();
    assert_eq!(extension_cache_stats(&mut engine, &ctx).entries, 1);
    assert!(matches!(
        first.tensor_read(),
        TensorRead::View(TensorView::F64(_))
    ));

    let second = einsum_with_engine(&mut engine, &[&lhs, &rhs], "kjt,ot->jkot").unwrap();

    assert_eq!(extension_cache_stats(&mut engine, &ctx).entries, 1);
    match second.tensor_read() {
        TensorRead::View(TensorView::F64(view)) => {
            assert_eq!(view.shape(), &[j, k, o, t]);
            assert_eq!(
                view.strides(),
                &[k as isize, 1, (k * j) as isize, (k * j * o) as isize]
            );
        }
        other => panic!("expected lazy f64 view from cached expansion, got {other:?}"),
    }
}

#[test]
fn nary_eager_einsum_expanded_standard_ops_preserve_backward() {
    let ctx = EagerRuntime::new().unwrap();
    let a = EagerTensor::requires_grad_in(
        Tensor::from_vec_col_major(vec![2, 3], vec![1.0_f64; 6]).unwrap(),
        ctx.clone(),
    )
    .unwrap();
    let b = EagerTensor::from_tensor_in(
        Tensor::from_vec_col_major(vec![3, 4], vec![1.0_f64; 12]).unwrap(),
        ctx.clone(),
    )
    .unwrap();
    let c = EagerTensor::from_tensor_in(
        Tensor::from_vec_col_major(vec![4, 5], vec![1.0_f64; 20]).unwrap(),
        ctx,
    )
    .unwrap();

    let out = einsum(&[&a, &b, &c], "ij,jk,kl->il").unwrap();
    let loss = out
        .runtime()
        .with_eager_session(&mut CpuEngine::serial(), |s| {
            s.reduce_sum(&out, Some(&[0, 1]))
        })
        .unwrap();
    let _ = loss.backward(&mut CpuEngine::serial()).unwrap();

    assert_eq!(
        a.grad().unwrap().unwrap().as_slice::<f64>().unwrap(),
        &[20.0; 6]
    );
}

#[test]
fn tracked_nary_einsum_gradients_match_expected_values() {
    let ctx = EagerRuntime::new().unwrap();
    let a = EagerTensor::requires_grad_in(
        Tensor::from_vec_col_major(vec![2, 3], vec![1.0_f64; 6]).unwrap(),
        ctx.clone(),
    )
    .unwrap();
    let b = EagerTensor::requires_grad_in(
        Tensor::from_vec_col_major(vec![3, 4], vec![1.0_f64; 12]).unwrap(),
        ctx.clone(),
    )
    .unwrap();
    let c = EagerTensor::requires_grad_in(
        Tensor::from_vec_col_major(vec![4, 5], vec![1.0_f64; 20]).unwrap(),
        ctx,
    )
    .unwrap();

    let out = einsum(&[&a, &b, &c], "ij,jk,kl->il").unwrap();
    assert_eq!(out.shape(), &[2, 5]);

    let loss = out
        .runtime()
        .with_eager_session(&mut CpuEngine::serial(), |s| {
            s.reduce_sum(&out, Some(&[0, 1]))
        })
        .unwrap();
    let _ = loss.backward(&mut CpuEngine::serial()).unwrap();

    assert_eq!(
        a.grad().unwrap().unwrap().as_slice::<f64>().unwrap(),
        &[20.0; 6]
    );
    assert_eq!(
        b.grad().unwrap().unwrap().as_slice::<f64>().unwrap(),
        &[10.0; 12]
    );
    assert_eq!(
        c.grad().unwrap().unwrap().as_slice::<f64>().unwrap(),
        &[6.0; 20]
    );
}

#[test]
fn eager_outer_product_can_use_untracked_backend_broadcast_multiply() {
    let ctx = EagerRuntime::new().unwrap();
    let lhs = EagerTensor::from_tensor_in(
        Tensor::from_vec_col_major(vec![2], vec![2.0_f64, 3.0]).unwrap(),
        ctx.clone(),
    )
    .unwrap();
    let rhs = EagerTensor::from_tensor_in(
        Tensor::from_vec_col_major(vec![3], vec![5.0_f64, 7.0, 11.0]).unwrap(),
        ctx.clone(),
    )
    .unwrap();

    let out = ctx
        .with_eager_session(&mut CpuEngine::serial(), |session| {
            backend_broadcast_multiply_untracked(session, &lhs, &[2, 3], &[0], &rhs, &[2, 3], &[1])
        })
        .unwrap()
        .expect("untracked CPU eager tensors should use backend broadcast multiply");

    assert_eq!(out.shape(), &[2, 3]);
    assert_eq!(
        out.value().unwrap().as_slice::<f64>().unwrap(),
        &[10.0, 15.0, 14.0, 21.0, 22.0, 33.0]
    );
    assert!(!out.tracks_grad());
}

#[test]
fn eager_outer_product_can_return_lazy_noncompact_output() {
    let ctx = EagerRuntime::new().unwrap();
    let k = 2;
    let j = 3;
    let o = 4;
    let t = 5;
    let lhs_data: Vec<f64> = (0..k * j * t).map(|idx| idx as f64 + 1.0).collect();
    let rhs_data: Vec<f64> = (0..o * t).map(|idx| idx as f64 + 101.0).collect();
    let lhs = EagerTensor::from_tensor_in(
        Tensor::from_vec_col_major(vec![k, j, t], lhs_data.clone()).unwrap(),
        ctx.clone(),
    )
    .unwrap();
    let rhs = EagerTensor::from_tensor_in(
        Tensor::from_vec_col_major(vec![o, t], rhs_data.clone()).unwrap(),
        ctx,
    )
    .unwrap();

    let out = einsum(&[&lhs, &rhs], "kjt,ot->jkot").unwrap();

    assert_eq!(out.shape(), &[j, k, o, t]);
    match out.tensor_read() {
        TensorRead::View(TensorView::F64(view)) => {
            assert_eq!(view.shape(), &[j, k, o, t]);
            assert_eq!(
                view.strides(),
                &[k as isize, 1, (k * j) as isize, (k * j * o) as isize]
            );
        }
        other => panic!("expected lazy f64 view, got {other:?}"),
    }

    let expected: Vec<f64> = (0..t)
        .flat_map(|tt| {
            let lhs_data = &lhs_data;
            let rhs_data = &rhs_data;
            (0..o).flat_map(move |oo| {
                (0..k).flat_map(move |kk| {
                    (0..j).map(move |jj| lhs_data[kk + k * jj + k * j * tt] * rhs_data[oo + o * tt])
                })
            })
        })
        .collect();
    assert_eq!(
        out.to_tensor(&mut CpuEngine::serial())
            .unwrap()
            .as_slice::<f64>()
            .unwrap(),
        expected.as_slice()
    );
}

/// #1946 F4: an untracked einsum inside `capture_trace` stays in the semantic
/// trace, so the VJP reaches the untracked input. The removed prototype gate
/// (`TENFERRO_EAGER_WHOLE_PROGRAM`) executed such calls through concrete tensors
/// and returned a fresh leaf, losing the trace.
#[test]
fn captured_untracked_einsum_keeps_its_vjp() {
    let runtime = EagerRuntime::new().unwrap();
    let x = EagerTensor::from_tensor_in(
        Tensor::from_vec_col_major(vec![2], vec![2.0_f64, 3.0]).unwrap(),
        runtime.clone(),
    )
    .unwrap();
    let seed = EagerTensor::from_tensor_in(
        Tensor::from_vec_col_major(vec![], vec![1.0_f64]).unwrap(),
        runtime.clone(),
    )
    .unwrap();
    let y = {
        let _capture = runtime.capture_trace();
        einsum(&[&x, &x, &x], "i,i,i->").unwrap()
    };
    let grad = runtime
        .vjp(&mut CpuEngine::serial(), &y, &x, &seed)
        .unwrap()
        .to_tensor(&mut CpuEngine::serial())
        .unwrap();
    // d/dx_i sum_j x_j^3 = 3 x_i^2.
    assert_eq!(grad.as_slice::<f64>().unwrap(), &[12.0, 27.0]);
}

/// #1946 F4: no environment variable may switch eager einsum onto an
/// alternate execution path.
#[test]
fn eager_einsum_has_no_environment_execution_gate() {
    for (name, source) in [
        ("eager_ad.rs", include_str!("../eager_ad.rs")),
        ("eager.rs", include_str!("../eager.rs")),
    ] {
        assert!(
            !source.contains("TENFERRO_EAGER_WHOLE_PROGRAM") && !source.contains("std::env::var"),
            "{name} must not read an environment switch for eager einsum execution"
        );
    }
}