use std::sync::Arc;
use computegraph::graph::GraphBuilder;
use computegraph::{OperationRole, ValueKey, ValueRef};
use num_complex::{Complex32, Complex64};
use tenferro_cpu::CpuBackend;
use tenferro_ops::input_key::TensorInputKey;
use tenferro_ops::std_tensor_op::StdTensorOp;
use tenferro_tensor::{Tensor, TypedTensor};
use tidu::{ADKey, PrimitiveBuilder, PrimitiveValue};
use crate::extension_runtime::ExtensionExecutor;
use super::{missing_tangent_base_key, zero_like_tensor, EagerPrimitiveBuilder};
#[test]
fn debug_summarizes_builder_without_tensor_payloads() {
let mut backend = CpuBackend::new();
let mut builder = EagerPrimitiveBuilder::new(&mut backend);
let id = builder.push_tensor(Arc::new(
Tensor::from_vec_col_major(vec![1], vec![1.0_f64]).unwrap(),
));
let _tensor = builder.tensor(id).unwrap();
let debug = format!("{builder:?}");
assert!(debug.contains("EagerPrimitiveBuilder"));
assert!(debug.contains("backend_type"));
assert!(debug.contains("has_extension_executor: false"));
assert!(debug.contains("results_len: 1"));
}
#[test]
fn debug_reports_extension_executor_presence() {
let mut backend = CpuBackend::new();
let mut executor = ExtensionExecutor::<CpuBackend>::new();
let builder = EagerPrimitiveBuilder::with_extension_executor(&mut backend, &mut executor);
let debug = format!("{builder:?}");
assert!(debug.contains("has_extension_executor: true"));
}
#[test]
fn new_builder_executes_standard_primitives_without_extension_executor() {
let mut backend = CpuBackend::new();
let mut builder = EagerPrimitiveBuilder::new(&mut backend);
let lhs = builder.push_tensor(Arc::new(
Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0]).unwrap(),
));
let rhs = builder.push_tensor(Arc::new(
Tensor::from_vec_col_major(vec![2], vec![3.0_f64, 4.0]).unwrap(),
));
let outputs = PrimitiveBuilder::add_primitive(
&mut builder,
StdTensorOp::Add,
vec![PrimitiveValue::Local(lhs), PrimitiveValue::Local(rhs)],
OperationRole::Primary,
);
assert_eq!(outputs.len(), 1);
assert_eq!(
builder
.tensor(outputs[0])
.unwrap()
.as_slice::<f64>()
.unwrap(),
&[4.0, 6.0]
);
}
#[test]
fn missing_tangent_external_uses_zero_like_primal_fallback() {
let mut backend = CpuBackend::new();
let mut builder = EagerPrimitiveBuilder::new(&mut backend);
let primal_input = TensorInputKey::User { id: 7 };
let primal_key = ValueKey::Input(primal_input.clone());
let tangent_key = ValueKey::Input(primal_input.tangent_of(3));
builder.external_data.insert(
primal_key,
Arc::new(Tensor::from_vec_col_major(vec![2], vec![5.0_f64, 7.0]).unwrap()),
);
let outputs = PrimitiveBuilder::add_primitive(
&mut builder,
StdTensorOp::Neg,
vec![PrimitiveValue::External(tangent_key.clone())],
OperationRole::Primary,
);
assert_eq!(outputs.len(), 1);
assert!(builder.external_data.contains_key(&tangent_key));
assert_eq!(
builder
.tensor(outputs[0])
.unwrap()
.as_slice::<f64>()
.unwrap(),
&[0.0, 0.0]
);
}
#[test]
fn missing_non_tangent_external_records_error_without_panicking() {
let mut backend = CpuBackend::new();
let mut builder = EagerPrimitiveBuilder::new(&mut backend);
let missing_key = ValueKey::Input(TensorInputKey::User { id: 19 });
let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
PrimitiveBuilder::add_primitive(
&mut builder,
StdTensorOp::Neg,
vec![PrimitiveValue::External(missing_key)],
OperationRole::Primary,
)
}));
assert!(result.is_ok());
assert!(builder.take_error().is_some());
}
#[test]
fn missing_tangent_base_key_accepts_only_input_tangent_keys() {
let primal_input = TensorInputKey::User { id: 11 };
let primal_key = ValueKey::Input(primal_input.clone());
assert_eq!(missing_tangent_base_key(&primal_key), None);
let tangent_key = ValueKey::Input(primal_input.tangent_of(5));
assert_eq!(missing_tangent_base_key(&tangent_key), Some(primal_key));
let mut graph = GraphBuilder::<StdTensorOp>::new();
let input = graph.add_input(TensorInputKey::User { id: 12 });
let output = graph.add_operation(
StdTensorOp::Neg,
vec![ValueRef::Local(input)],
OperationRole::Primary,
)[0];
let derived_key = graph.global_key(output).clone();
assert_eq!(missing_tangent_base_key(&derived_key), None);
}
#[test]
fn zero_like_tensor_covers_all_dtypes() {
assert_zero_like_matches(Tensor::F32(
TypedTensor::from_vec_col_major(vec![2], vec![1.0_f32, -2.0]).unwrap(),
));
assert_zero_like_matches(Tensor::F64(
TypedTensor::from_vec_col_major(vec![2], vec![1.0_f64, -2.0]).unwrap(),
));
assert_zero_like_matches(Tensor::I32(
TypedTensor::from_vec_col_major(vec![2], vec![1_i32, -2]).unwrap(),
));
assert_zero_like_matches(Tensor::I64(
TypedTensor::from_vec_col_major(vec![2], vec![1_i64, -2]).unwrap(),
));
assert_zero_like_matches(Tensor::Bool(
TypedTensor::from_vec_col_major(vec![2], vec![true, false]).unwrap(),
));
assert_zero_like_matches(Tensor::C32(
TypedTensor::from_vec_col_major(
vec![2],
vec![Complex32::new(1.0, 2.0), Complex32::new(-3.0, 4.0)],
)
.unwrap(),
));
assert_zero_like_matches(Tensor::C64(
TypedTensor::from_vec_col_major(
vec![2],
vec![Complex64::new(1.0, 2.0), Complex64::new(-3.0, 4.0)],
)
.unwrap(),
));
}
fn assert_zero_like_matches(input: Tensor) {
let shape = input.shape().to_vec();
let mut backend = CpuBackend::new();
let zero = zero_like_tensor(&input, &mut backend).unwrap();
assert_eq!(zero.shape(), shape.as_slice());
match zero {
Tensor::F32(tensor) => assert_eq!(tensor.as_slice().unwrap(), &[0.0_f32, 0.0]),
Tensor::F64(tensor) => assert_eq!(tensor.as_slice().unwrap(), &[0.0_f64, 0.0]),
Tensor::I32(tensor) => assert_eq!(tensor.as_slice().unwrap(), &[0_i32, 0]),
Tensor::I64(tensor) => assert_eq!(tensor.as_slice().unwrap(), &[0_i64, 0]),
Tensor::Bool(tensor) => assert_eq!(tensor.as_slice().unwrap(), &[false, false]),
Tensor::C32(tensor) => assert_eq!(
tensor.as_slice().unwrap(),
&[Complex32::new(0.0, 0.0), Complex32::new(0.0, 0.0)]
),
Tensor::C64(tensor) => assert_eq!(
tensor.as_slice().unwrap(),
&[Complex64::new(0.0, 0.0), Complex64::new(0.0, 0.0)]
),
}
}