use std::borrow::Cow;
use std::cell::{Cell, RefCell};
use std::cmp::Reverse;
use std::collections::HashMap;
use std::env;
use std::fmt;
use std::marker::PhantomData;
use std::mem::{size_of, size_of_val};
use std::rc::Rc;
#[cfg(test)]
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Mutex, MutexGuard, OnceLock, Weak};
use std::time::{Duration, Instant};
use lru::LruCache;
use num_complex::Complex64;
use crate::extension::{
validate_eager_extension_target, EagerExtensionBackendKind, EagerExtensionTarget,
};
#[cfg(test)]
use computegraph::graph::Graph;
use computegraph::ValueKey;
#[cfg(test)]
use computegraph::ValueRef;
#[cfg(test)]
use tenferro_ops::input_key::TensorInputKey;
use tenferro_ops::{std_tensor_op::StdTensorOp, SymDim, TensorMeta};
use tenferro_runtime::ad_support::{
analyze_deferred_semantic_trace, compile_ad_source, ones_tensor, RetainedValue,
};
use tenferro_runtime::program::{ProgramValueMetadata, SemanticFingerprint, SemanticProgram};
use tenferro_runtime::{
CompiledGraph, EngineId, EngineTarget, ErrorPhase, ExecutionContextIdentity, ExecutionEngine,
ExtensionModule, GraphCompiler, Prepared, Runtime, RuntimeConfigError, RuntimeEpoch,
TracedTensor,
};
use tenferro_tensor::Engine;
#[cfg(test)]
use tenferro_tensor::TypedTensor;
use tenferro_tensor::{
AllocationGroup, CacheStats, CompareDir, DType, DescriptorSlot, DotGeneralConfig, GatherConfig,
GroupError, IntoShapeVec, PadConfig, ScatterConfig, SliceConfig, Tensor, TensorRead,
TensorScalar, TensorValue, TensorView,
};
use crate::eager_backend::{
cpu_host_owned_read, engine_kind_registration, unregisterable_engine_error,
};
#[cfg(test)]
use crate::eager_exec::exec_standard_op_on_tensor_reads_in_session;
use crate::eager_exec::{eager_input_promotion_plan, exec_extension_op_on_tensor_reads};
use crate::error::{ContextId, Error, Result};
use crate::metadata::tensor_meta_from_tensor;
use crate::semantic_extension::SemanticExtensionRuleSet;
use crate::traced::{derivative_trace_from_frozen_program, next_input_key};
use crate::transform_cache::{AdTransformCache, AdTransformCacheLimits};
use crate::AdContext;
pub(crate) type GradSlot = Arc<Mutex<Option<Arc<AdValueRecord>>>>;
pub(crate) type WeakGradSlot = Weak<Mutex<Option<Arc<AdValueRecord>>>>;
mod composite;
mod residuals;
pub(crate) use residuals::{finish_residuals, EagerTrace};
#[derive(Debug, Default, Clone)]
struct EagerOpProfileEntry {
calls: usize,
total_time: Duration,
}
thread_local! {
static EAGER_OP_PROFILE_STATE: RefCell<HashMap<&'static str, EagerOpProfileEntry>> =
RefCell::new(HashMap::new());
static EAGER_NO_GRAD_DEPTH: Cell<usize> = const { Cell::new(0) };
static EAGER_CAPTURE_DEPTH: Cell<usize> = const { Cell::new(0) };
#[cfg(test)]
static EAGER_OP_PROFILE_ENABLED_OVERRIDE: RefCell<Option<bool>> = const { RefCell::new(None) };
#[cfg(test)]
static EAGER_OP_PROFILE_PRINT_EVERY_OVERRIDE: RefCell<Option<Option<usize>>> = const { RefCell::new(None) };
#[cfg(test)]
static EAGER_SEMANTIC_VJP_ENABLED_OVERRIDE: RefCell<Option<bool>> = const { RefCell::new(None) };
}
#[cfg(test)]
pub(crate) static EAGER_SEMANTIC_VJP_EXECUTIONS: AtomicUsize = AtomicUsize::new(0);
pub(crate) fn eager_grad_recording_enabled() -> bool {
EAGER_NO_GRAD_DEPTH.with(|depth| depth.get() == 0)
}
pub(crate) fn eager_capture_active() -> bool {
EAGER_CAPTURE_DEPTH.with(|depth| depth.get() > 0)
}
fn eager_semantic_vjp_enabled() -> bool {
#[cfg(test)]
if let Some(value) = EAGER_SEMANTIC_VJP_ENABLED_OVERRIDE.with(|state| *state.borrow()) {
return value;
}
static ENABLED: OnceLock<bool> = OnceLock::new();
*ENABLED.get_or_init(|| env::var("TENFERRO_EAGER_SEMANTIC_VJP").map_or(true, |v| v != "0"))
}
#[derive(Debug)]
pub struct EagerNoGradGuard {
active: bool,
_not_send: PhantomData<Rc<()>>,
}
impl Drop for EagerNoGradGuard {
fn drop(&mut self) {
if !self.active {
return;
}
EAGER_NO_GRAD_DEPTH.with(|depth| {
depth.set(depth.get().saturating_sub(1));
});
self.active = false;
}
}
#[derive(Debug)]
pub struct EagerTraceCaptureGuard {
active: bool,
_not_send: PhantomData<Rc<()>>,
}
impl Drop for EagerTraceCaptureGuard {
fn drop(&mut self) {
if !self.active {
return;
}
EAGER_CAPTURE_DEPTH.with(|depth| {
depth.set(depth.get().saturating_sub(1));
});
self.active = false;
}
}
pub(crate) fn eager_op_profile_enabled() -> bool {
#[cfg(test)]
if let Some(value) = EAGER_OP_PROFILE_ENABLED_OVERRIDE.with(|state| *state.borrow()) {
return value;
}
static ENABLED: OnceLock<bool> = OnceLock::new();
*ENABLED.get_or_init(|| env::var("TENFERRO_PROFILE_EAGER_OP_AGG").is_ok())
}
pub(crate) fn eager_op_profile_start() -> Option<Instant> {
eager_op_profile_enabled().then(Instant::now)
}
pub(crate) fn record_eager_op_profile(section: &'static str, elapsed: Duration) {
if !eager_op_profile_enabled() {
return;
}
EAGER_OP_PROFILE_STATE.with(|state| {
let mut state = state.borrow_mut();
let entry = state.entry(section).or_default();
entry.calls += 1;
entry.total_time += elapsed;
});
}
pub(crate) fn profile_eager_op_section<T>(section: &'static str, f: impl FnOnce() -> T) -> T {
if !eager_op_profile_enabled() {
return f();
}
let started = Instant::now();
let result = f();
record_eager_op_profile(section, started.elapsed());
result
}
pub(crate) fn maybe_print_eager_op_profile() {
if !eager_op_profile_enabled() {
return;
}
let Some(print_every) = eager_op_profile_print_every() else {
return;
};
if print_every == 0 {
return;
}
let should_print = EAGER_OP_PROFILE_STATE.with(|state| {
state
.borrow()
.get("nary_op.total")
.is_some_and(|entry| entry.calls % print_every == 0)
});
if should_print {
print_and_reset_eager_op_profile();
}
}
fn eager_op_profile_print_every() -> Option<usize> {
#[cfg(test)]
if let Some(value) = EAGER_OP_PROFILE_PRINT_EVERY_OVERRIDE.with(|state| *state.borrow()) {
return value;
}
env::var("TENFERRO_PROFILE_EAGER_OP_PRINT_EVERY")
.ok()?
.parse()
.ok()
}
pub(crate) fn print_and_reset_eager_op_profile() {
EAGER_OP_PROFILE_STATE.with(|state| {
let mut entries: Vec<_> = state
.borrow()
.iter()
.map(|(section, entry)| (*section, entry.clone()))
.collect();
state.borrow_mut().clear();
entries.sort_by_key(|(_, entry)| Reverse(entry.total_time));
eprintln!("=== tenferro eager op profile ===");
for (section, entry) in entries {
let Some(per_call_us) = eager_op_profile_per_call_us(&entry) else {
continue;
};
eprintln!(
"{section}: calls={} total={:.6}ms per_call={:.3}us",
entry.calls,
entry.total_time.as_secs_f64() * 1.0e3,
per_call_us,
);
}
});
}
fn eager_op_profile_per_call_us(entry: &EagerOpProfileEntry) -> Option<f64> {
(entry.calls != 0).then(|| entry.total_time.as_secs_f64() * 1.0e6 / entry.calls as f64)
}
fn runtime_config_error(op: &'static str, source: RuntimeConfigError) -> Error {
Error::runtime_state_source(op, ErrorPhase::Execution, source)
}
fn runtime_state_source<E>(op: &'static str, source: E) -> Error
where
E: std::error::Error + Send + Sync + 'static,
{
Error::runtime_state_source(op, ErrorPhase::Execution, source)
}
pub(crate) fn run_prepared_on(
runtime: &Runtime,
engine: &mut dyn ExecutionEngine,
prepared: &Prepared,
inputs: &[TensorRead<'_>],
) -> Result<Vec<Tensor>> {
let outputs = runtime.run(&mut [&mut *engine], prepared, inputs)?;
outputs
.into_iter()
.map(|output| output.into_tensor(&mut *engine).map_err(Error::from))
.collect()
}
pub(crate) fn execute_prepared_extension_on(
engine: &mut dyn ExecutionEngine,
executor: &dyn tenferro_runtime::PreparedOperationExecutor,
inputs: &[TensorRead<'_>],
) -> Result<Vec<Tensor>> {
with_lent_extension_caches(engine, |engine, caches| {
executor.execute_in_session(engine, caches, inputs)
})
}
pub(crate) fn with_lent_extension_caches<R>(
engine: &mut dyn ExecutionEngine,
callback: impl FnOnce(&mut dyn ExecutionEngine, &mut tenferro_runtime::ExtensionCacheStore) -> R,
) -> R {
let mut caches = engine.take_extension_caches();
let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
callback(&mut *engine, &mut caches)
}));
engine.restore_extension_caches(caches);
match result {
Ok(value) => value,
Err(payload) => {
engine.mark_broken();
std::panic::resume_unwind(payload)
}
}
}
pub(crate) fn eager_extension_backend_kind(
target: &EngineTarget,
) -> Option<EagerExtensionBackendKind> {
if target.context_identity() == ExecutionContextIdentity::of::<tenferro_cpu::CpuEngine>() {
return Some(EagerExtensionBackendKind::Cpu);
}
#[cfg(feature = "cuda")]
if target.context_identity() == ExecutionContextIdentity::of::<tenferro_gpu::cuda::CudaEngine>()
{
return Some(EagerExtensionBackendKind::Cuda);
}
None
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct EagerRuntimeCacheStats {
pub ad_transforms: CacheStats,
pub prepared_derivatives: CacheStats,
}
#[cfg(test)]
pub(crate) struct EagerGraphExecution {
pub(crate) outputs: Vec<Tensor>,
}
#[derive(Debug)]
pub struct ValueGuard<'a> {
view: TensorView<'a>,
}
impl<'a> ValueGuard<'a> {
pub fn dtype(&self) -> DType {
self.view.dtype()
}
pub fn shape(&self) -> &[usize] {
self.view.shape()
}
pub fn as_tensor_view(&self) -> &TensorView<'_> {
&self.view
}
pub fn as_slice<T: TensorScalar>(&self) -> tenferro_tensor::Result<&'a [T]> {
self.view.as_slice()
}
fn duplicate_host_tensor(&self) -> tenferro_tensor::Result<Tensor> {
match &self.view {
TensorView::F32(view) => {
<f32 as TensorScalar>::into_tensor(view.shape().to_vec(), view.as_slice()?.to_vec())
}
TensorView::F64(view) => {
<f64 as TensorScalar>::into_tensor(view.shape().to_vec(), view.as_slice()?.to_vec())
}
TensorView::I32(view) => {
<i32 as TensorScalar>::into_tensor(view.shape().to_vec(), view.as_slice()?.to_vec())
}
TensorView::I64(view) => {
<i64 as TensorScalar>::into_tensor(view.shape().to_vec(), view.as_slice()?.to_vec())
}
TensorView::Bool(view) => <bool as TensorScalar>::into_tensor(
view.shape().to_vec(),
view.as_slice()?.to_vec(),
),
TensorView::C32(view) => <num_complex::Complex32 as TensorScalar>::into_tensor(
view.shape().to_vec(),
view.as_slice()?.to_vec(),
),
TensorView::C64(view) => <num_complex::Complex64 as TensorScalar>::into_tensor(
view.shape().to_vec(),
view.as_slice()?.to_vec(),
),
}
}
}
#[derive(Clone, Debug)]
pub struct GradientValue {
record: Arc<AdValueRecord>,
}
impl GradientValue {
pub fn dtype(&self) -> DType {
self.record.dtype()
}
pub fn shape(&self) -> &[usize] {
self.record.shape()
}
pub fn value(&self) -> Result<ValueGuard<'_>> {
self.record.value("GradientValue::value")
}
pub fn tensor_read(&self) -> Result<TensorRead<'_>> {
self.record.tensor_read("GradientValue::tensor_read")
}
pub fn as_slice<T: TensorScalar>(&self) -> tenferro_tensor::Result<&[T]> {
self.record
.value("GradientValue::as_slice")
.map_err(|error| {
tenferro_tensor::Error::runtime_state_source("GradientValue::as_slice", error)
})?
.as_slice()
}
pub fn to_tensor(&self, engine: &mut dyn Engine) -> Result<Tensor> {
let value = self
.record
.value("GradientValue::to_tensor")
.map_err(|error| {
Error::runtime_state_source(
"GradientValue::to_tensor",
ErrorPhase::Execution,
error,
)
})?;
match value.duplicate_host_tensor() {
Ok(tensor) => Ok(tensor),
Err(_) => {
let read = self.record.tensor_read("GradientValue::to_tensor")?;
engine.to_contiguous_read(read).map_err(Error::from)
}
}
}
}
#[derive(Debug)]
pub struct Gradients {
group: AllocationGroup,
slots: HashMap<ValueKey<StdTensorOp>, DescriptorSlot>,
}
impl Gradients {
fn from_tensors(tensors: HashMap<ValueKey<StdTensorOp>, Tensor>) -> Result<Self> {
let (keys, values): (Vec<_>, Vec<_>) = tensors.into_iter().unzip();
let (group, bindings) = AllocationGroup::from_tensors(values).map_err(|error| {
Error::runtime_state_source("Gradients::from_tensors", ErrorPhase::Execution, error)
})?;
let slots = keys.into_iter().zip(bindings).collect();
Ok(Self { group, slots })
}
pub fn len(&self) -> usize {
self.slots.len()
}
pub fn is_empty(&self) -> bool {
self.slots.is_empty()
}
pub fn grad(&self, key: &ValueKey<StdTensorOp>) -> Option<TensorView<'_>> {
let slot = self.slots.get(key).copied()?;
let mut reads = self.group.read_views(std::slice::from_ref(&slot)).ok()?;
match reads.pop()? {
TensorRead::View(view) => Some(view),
TensorRead::Tensor(_) => None,
}
}
pub fn take_grad(
&mut self,
key: &ValueKey<StdTensorOp>,
) -> tenferro_tensor::Result<Option<Tensor>> {
let Some(&slot) = self.slots.get(key) else {
return Ok(None);
};
let tensor = self.group.take_tensor(slot).map_err(|error| {
tenferro_tensor::Error::runtime_state_source("Gradients::take_grad", error)
})?;
self.slots.remove(key);
Ok(Some(tensor))
}
}
#[derive(Debug)]
pub enum IntoValueError<H> {
NotUnique(H),
Extract { value: H, error: GroupError },
}
impl<H> std::fmt::Display for IntoValueError<H> {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::NotUnique(_) => formatter.write_str("eager value is retained by another handle"),
Self::Extract { error, .. } => {
write!(formatter, "eager value extraction failed: {error}")
}
}
}
}
impl<H: std::fmt::Debug + Send + Sync + 'static> std::error::Error for IntoValueError<H> {}
#[allow(clippy::large_enum_variant)]
#[derive(Debug)]
enum RetentionContainer {
Pooled {
group: AllocationGroup,
slot: DescriptorSlot,
},
CallerOwned {
tensor: Box<Tensor>,
},
Owned { tensor: Tensor },
}
#[derive(Debug)]
pub(crate) struct AdValueRecord {
container: Arc<RetentionContainer>,
dtype: DType,
shape: Box<[usize]>,
}
impl AdValueRecord {
fn from_group(
group: AllocationGroup,
slot: DescriptorSlot,
dtype: DType,
shape: Vec<usize>,
) -> Arc<Self> {
Arc::new(Self {
container: Arc::new(RetentionContainer::Pooled { group, slot }),
dtype,
shape: shape.into_boxed_slice(),
})
}
fn from_tensor(tensor: Tensor, op: &'static str) -> Result<Arc<Self>> {
let dtype = tensor.dtype();
let shape = tensor.shape().to_vec();
if matches!(dtype, DType::External(_)) {
return Ok(Arc::new(Self {
container: Arc::new(RetentionContainer::CallerOwned {
tensor: Box::new(tensor),
}),
dtype,
shape: shape.into_boxed_slice(),
}));
}
let (group, bindings) = AllocationGroup::from_tensors(vec![tensor])
.map_err(|error| Error::runtime_state_source(op, ErrorPhase::Execution, error))?;
let slot = bindings.first().copied().ok_or_else(|| {
Error::runtime_state(op, ErrorPhase::Execution, "empty allocation-group binding")
})?;
Ok(Self::from_group(group, slot, dtype, shape))
}
fn from_untracked_tensor(tensor: Tensor, op: &'static str) -> Result<Arc<Self>> {
if matches!(tensor.dtype(), DType::External(_)) {
return Self::from_tensor(tensor, op);
}
let dtype = tensor.dtype();
let shape = tensor.shape().to_vec().into_boxed_slice();
Ok(Arc::new(Self {
container: Arc::new(RetentionContainer::Owned { tensor }),
dtype,
shape,
}))
}
fn tensor_read(&self, op: &'static str) -> Result<TensorRead<'_>> {
match self.container.as_ref() {
RetentionContainer::Pooled { group, slot } => group.read_view(*slot).map_err(|error| {
Error::runtime_state_source(op, ErrorPhase::Execution, error)
}),
RetentionContainer::CallerOwned { tensor } => Ok(TensorRead::from_tensor(tensor)),
RetentionContainer::Owned { tensor } => Ok(TensorRead::from_tensor(tensor)),
}
}
fn value(&self, op: &'static str) -> Result<ValueGuard<'_>> {
if let RetentionContainer::Owned { tensor } = self.container.as_ref() {
return Ok(ValueGuard {
view: TensorRead::from_tensor(tensor).tensor_view(),
});
}
match self.tensor_read(op)? {
TensorRead::View(view) => Ok(ValueGuard { view }),
TensorRead::Tensor(_) => Err(Error::runtime_state(
op,
ErrorPhase::Execution,
"allocation-group value did not produce a borrowed descriptor view",
)),
}
}
fn dtype(&self) -> DType {
self.dtype
}
fn shape(&self) -> &[usize] {
&self.shape
}
}
pub struct EagerRuntime {
id: ContextId,
runtime: Runtime,
extension_install_lock: Mutex<()>,
semantic_extension_rules: SemanticExtensionRuleSet,
grad_slots: Mutex<HashMap<ValueKey<StdTensorOp>, WeakGradSlot>>,
value_records: Mutex<HashMap<ValueKey<StdTensorOp>, Weak<EagerTensorRecord>>>,
ad_transform_cache: Arc<AdTransformCache>,
prepared_derivative_cache: Mutex<PreparedDerivativeCache>,
}
pub struct EagerSession<'a> {
runtime: &'a Arc<EagerRuntime>,
engine: &'a mut dyn ExecutionEngine,
}
impl fmt::Debug for EagerSession<'_> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("EagerSession")
.field("runtime_id", &self.runtime.id())
.finish_non_exhaustive()
}
}
impl EagerSession<'_> {
pub fn neg(&mut self, input: &EagerTensor) -> Result<EagerTensor> {
self.run_unary(input, StdTensorOp::Neg)
}
pub fn exp(&mut self, input: &EagerTensor) -> Result<EagerTensor> {
self.run_unary(input, StdTensorOp::Exp)
}
pub fn abs(&mut self, input: &EagerTensor) -> Result<EagerTensor> {
self.run_unary(input, StdTensorOp::Abs)
}
pub fn conj(&mut self, input: &EagerTensor) -> Result<EagerTensor> {
self.run_unary(input, StdTensorOp::Conj)
}
pub fn sign(&mut self, input: &EagerTensor) -> Result<EagerTensor> {
self.run_unary(input, StdTensorOp::Sign)
}
pub fn log(&mut self, input: &EagerTensor) -> Result<EagerTensor> {
self.run_unary(input, StdTensorOp::Log)
}
pub fn sqrt(&mut self, input: &EagerTensor) -> Result<EagerTensor> {
self.run_unary(input, StdTensorOp::Sqrt)
}
pub fn rsqrt(&mut self, input: &EagerTensor) -> Result<EagerTensor> {
self.run_unary(input, StdTensorOp::Rsqrt)
}
pub fn sin(&mut self, input: &EagerTensor) -> Result<EagerTensor> {
self.run_unary(input, StdTensorOp::Sin)
}
pub fn cos(&mut self, input: &EagerTensor) -> Result<EagerTensor> {
self.run_unary(input, StdTensorOp::Cos)
}
pub fn tanh(&mut self, input: &EagerTensor) -> Result<EagerTensor> {
self.run_unary(input, StdTensorOp::Tanh)
}
pub fn expm1(&mut self, input: &EagerTensor) -> Result<EagerTensor> {
self.run_unary(input, StdTensorOp::Expm1)
}
pub fn log1p(&mut self, input: &EagerTensor) -> Result<EagerTensor> {
self.run_unary(input, StdTensorOp::Log1p)
}
pub fn erf(&mut self, input: &EagerTensor) -> Result<EagerTensor> {
self.run_unary(input, StdTensorOp::Erf)
}
pub fn convert(&mut self, input: &EagerTensor, to: DType) -> Result<EagerTensor> {
self.ensure_runtime(input)?;
tenferro_tensor::validate::validate_convert_dtype(
"EagerTensor::convert",
input.dtype(),
to,
)
.map_err(Error::TensorRuntime)?;
self.cast(input, to)
}
pub fn cast(&mut self, input: &EagerTensor, to: DType) -> Result<EagerTensor> {
self.run_unary(
input,
StdTensorOp::Convert {
from: input.dtype(),
to,
},
)
}
pub fn transpose(&mut self, input: &EagerTensor, perm: &[usize]) -> Result<EagerTensor> {
self.ensure_runtime(input)?;
input.transpose_in_session(perm, &mut *self.engine)
}
pub fn reshape(
&mut self,
input: &EagerTensor,
shape: impl IntoShapeVec,
) -> Result<EagerTensor> {
self.ensure_runtime(input)?;
input.reshape_in_session(&shape.into_shape_vec(), &mut *self.engine)
}
pub fn slice(&mut self, input: &EagerTensor, config: SliceConfig) -> Result<EagerTensor> {
self.ensure_runtime(input)?;
input.slice_in_session(config, &mut *self.engine)
}
pub fn broadcast_in_dim(
&mut self,
input: &EagerTensor,
shape: &[usize],
dims: &[usize],
) -> Result<EagerTensor> {
self.ensure_runtime(input)?;
input.broadcast_in_dim_in_session(shape, dims, &mut *self.engine)
}
pub fn tril(&mut self, input: &EagerTensor, k: i64) -> Result<EagerTensor> {
self.run_unary(input, StdTensorOp::Tril { k })
}
pub fn triu(&mut self, input: &EagerTensor, k: i64) -> Result<EagerTensor> {
self.run_unary(input, StdTensorOp::Triu { k })
}
pub fn pad(&mut self, input: &EagerTensor, config: PadConfig) -> Result<EagerTensor> {
self.run_unary(input, StdTensorOp::Pad(config))
}
pub fn reverse(&mut self, input: &EagerTensor, axes: &[usize]) -> Result<EagerTensor> {
self.ensure_runtime(input)?;
crate::eager_ops::validate_eager_axes("EagerSession::reverse", input.shape().len(), axes)?;
self.run_unary(
input,
StdTensorOp::Reverse {
axes: axes.to_vec(),
},
)
}
pub fn dynamic_slice(
&mut self,
input: &EagerTensor,
starts: &EagerTensor,
sizes: &[usize],
) -> Result<EagerTensor> {
self.ensure_runtime(input)?;
self.ensure_runtime(starts)?;
EagerTensor::nary_op_in_session(
&[input, starts],
StdTensorOp::DynamicSlice {
slice_sizes: sizes.to_vec(),
},
&mut *self.engine,
)
}
pub fn gather(
&mut self,
input: &EagerTensor,
indices: &EagerTensor,
config: GatherConfig,
) -> Result<EagerTensor> {
self.ensure_runtime(input)?;
self.ensure_runtime(indices)?;
EagerTensor::nary_op_in_session(
&[input, indices],
StdTensorOp::Gather(config),
&mut *self.engine,
)
}
pub fn concatenate(&mut self, inputs: &[&EagerTensor], axis: usize) -> Result<EagerTensor> {
for input in inputs {
self.ensure_runtime(input)?;
}
EagerTensor::nary_op_in_session(
inputs,
StdTensorOp::Concatenate {
axis,
input_count: inputs.len(),
},
&mut *self.engine,
)
}
pub fn scatter(
&mut self,
input: &EagerTensor,
indices: &EagerTensor,
updates: &EagerTensor,
config: ScatterConfig,
) -> Result<EagerTensor> {
self.ensure_runtime(input)?;
self.ensure_runtime(indices)?;
self.ensure_runtime(updates)?;
EagerTensor::nary_op_in_session(
&[input, indices, updates],
StdTensorOp::Scatter(config),
&mut *self.engine,
)
}
pub fn extract_diag(
&mut self,
input: &EagerTensor,
axis_a: usize,
axis_b: usize,
) -> Result<EagerTensor> {
self.ensure_runtime(input)?;
self.run_unary(input, StdTensorOp::ExtractDiag { axis_a, axis_b })
}
pub fn embed_diag(
&mut self,
input: &EagerTensor,
axis_a: usize,
axis_b: usize,
) -> Result<EagerTensor> {
self.ensure_runtime(input)?;
self.run_unary(input, StdTensorOp::EmbedDiag { axis_a, axis_b })
}
pub fn reduce_sum(
&mut self,
input: &EagerTensor,
axes: Option<&[usize]>,
) -> Result<EagerTensor> {
self.ensure_runtime(input)?;
input.reduce_sum_in_session(axes, &mut *self.engine)
}
pub fn reduce_sum_squares(
&mut self,
input: &EagerTensor,
axes: Option<&[usize]>,
) -> Result<EagerTensor> {
self.ensure_runtime(input)?;
let axes = axes.map_or_else(|| (0..input.shape().len()).collect(), <[usize]>::to_vec);
crate::eager_ops::validate_eager_axes(
"EagerSession::reduce_sum_squares",
input.shape().len(),
&axes,
)?;
self.run_unary(input, StdTensorOp::ReduceSumSquares { axes })
}
pub fn reduce_prod(
&mut self,
input: &EagerTensor,
axes: Option<&[usize]>,
) -> Result<EagerTensor> {
self.ensure_runtime(input)?;
let axes = axes.map_or_else(|| (0..input.shape().len()).collect(), <[usize]>::to_vec);
crate::eager_ops::validate_eager_axes(
"EagerSession::reduce_prod",
input.shape().len(),
&axes,
)?;
self.run_unary(input, StdTensorOp::ReduceProd { axes })
}
pub fn reduce_max(
&mut self,
input: &EagerTensor,
axes: Option<&[usize]>,
) -> Result<EagerTensor> {
self.ensure_runtime(input)?;
let axes = axes.map_or_else(|| (0..input.shape().len()).collect(), <[usize]>::to_vec);
crate::eager_ops::validate_eager_axes(
"EagerSession::reduce_max",
input.shape().len(),
&axes,
)?;
self.run_unary(input, StdTensorOp::ReduceMax { axes })
}
pub fn reduce_min(
&mut self,
input: &EagerTensor,
axes: Option<&[usize]>,
) -> Result<EagerTensor> {
self.ensure_runtime(input)?;
let axes = axes.map_or_else(|| (0..input.shape().len()).collect(), <[usize]>::to_vec);
crate::eager_ops::validate_eager_axes(
"EagerSession::reduce_min",
input.shape().len(),
&axes,
)?;
self.run_unary(input, StdTensorOp::ReduceMin { axes })
}
pub fn duplicate_value(&mut self, input: &EagerTensor) -> Result<Tensor> {
self.ensure_runtime(input)?;
input.duplicate_value_in_session(&mut *self.engine)
}
pub fn constant_from(&mut self, tensor: Tensor) -> Result<EagerTensor> {
EagerTensor::new_leaf_in_session(Arc::clone(self.runtime), tensor, false, &mut *self.engine)
}
pub fn constant_from_host(&mut self, tensor: Tensor) -> Result<EagerTensor> {
let uploaded = self
.engine
.upload_host_tensor(TensorRead::from_tensor(&tensor))
.map_err(Error::from)?;
self.constant_from(uploaded)
}
pub fn variable_from(&mut self, tensor: Tensor) -> Result<EagerTensor> {
EagerTensor::new_leaf_in_session(Arc::clone(self.runtime), tensor, true, &mut *self.engine)
}
pub fn add(&mut self, lhs: &EagerTensor, rhs: &EagerTensor) -> Result<EagerTensor> {
self.run_binary("add", lhs, rhs, StdTensorOp::Add)
}
pub fn sub(&mut self, lhs: &EagerTensor, rhs: &EagerTensor) -> Result<EagerTensor> {
self.run_binary("sub", lhs, rhs, StdTensorOp::Sub)
}
pub fn mul(&mut self, lhs: &EagerTensor, rhs: &EagerTensor) -> Result<EagerTensor> {
self.run_binary("mul", lhs, rhs, StdTensorOp::Mul)
}
pub fn div(&mut self, lhs: &EagerTensor, rhs: &EagerTensor) -> Result<EagerTensor> {
self.run_binary("div", lhs, rhs, StdTensorOp::Div)
}
pub fn rem(&mut self, lhs: &EagerTensor, rhs: &EagerTensor) -> Result<EagerTensor> {
self.run_binary("rem", lhs, rhs, StdTensorOp::Rem)
}
pub fn pow(&mut self, lhs: &EagerTensor, rhs: &EagerTensor) -> Result<EagerTensor> {
self.run_binary("pow", lhs, rhs, StdTensorOp::Pow)
}
pub fn maximum(&mut self, lhs: &EagerTensor, rhs: &EagerTensor) -> Result<EagerTensor> {
self.run_binary("maximum", lhs, rhs, StdTensorOp::Maximum)
}
pub fn minimum(&mut self, lhs: &EagerTensor, rhs: &EagerTensor) -> Result<EagerTensor> {
self.run_binary("minimum", lhs, rhs, StdTensorOp::Minimum)
}
pub fn compare(
&mut self,
lhs: &EagerTensor,
rhs: &EagerTensor,
dir: CompareDir,
) -> Result<EagerTensor> {
self.run_binary("compare", lhs, rhs, StdTensorOp::Compare(dir))
}
pub fn where_select(
&mut self,
condition: &EagerTensor,
on_true: &EagerTensor,
on_false: &EagerTensor,
) -> Result<EagerTensor> {
self.run_ternary(
"where_select",
condition,
on_true,
on_false,
StdTensorOp::Select,
)
}
pub fn select(
&mut self,
condition: &EagerTensor,
on_true: &EagerTensor,
on_false: &EagerTensor,
) -> Result<EagerTensor> {
self.where_select(condition, on_true, on_false)
}
pub fn clamp(
&mut self,
input: &EagerTensor,
lower: &EagerTensor,
upper: &EagerTensor,
) -> Result<EagerTensor> {
self.run_ternary("clamp", input, lower, upper, StdTensorOp::Clamp)
}
pub fn dot_general(
&mut self,
lhs: &EagerTensor,
rhs: &EagerTensor,
config: DotGeneralConfig,
) -> Result<EagerTensor> {
self.ensure_runtime(lhs)?;
self.ensure_runtime(rhs)?;
config
.validate_dims_with_ranks(lhs.shape().len(), rhs.shape().len())
.map_err(Error::TensorRuntime)?;
EagerTensor::nary_op_in_session(
&[lhs, rhs],
StdTensorOp::DotGeneral { config },
&mut *self.engine,
)
}
pub fn scale_real(&mut self, input: &EagerTensor, factor: f64) -> Result<EagerTensor> {
self.ensure_runtime(input)?;
let scalar = tenferro_runtime::scale::real_scale_scalar(input.dtype(), factor)?;
let scalar = self.constant_from_host(scalar)?;
self.mul(input, &scalar)
}
pub fn scale_complex(&mut self, input: &EagerTensor, factor: Complex64) -> Result<EagerTensor> {
self.ensure_runtime(input)?;
let scalar = tenferro_runtime::scale::complex_scale_scalar(input.dtype(), factor)?;
let scalar = self.constant_from_host(scalar)?;
self.mul(input, &scalar)
}
pub fn matmul(&mut self, lhs: &EagerTensor, rhs: &EagerTensor) -> Result<EagerTensor> {
self.ensure_runtime(lhs)?;
self.ensure_runtime(rhs)?;
let lhs_shape = lhs.shape();
let rhs_shape = rhs.shape();
if lhs_shape.len() != 2 {
return Err(tenferro_tensor::Error::rank_mismatch("matmul", 2, lhs_shape.len()).into());
}
if rhs_shape.len() != 2 {
return Err(tenferro_tensor::Error::rank_mismatch("matmul", 2, rhs_shape.len()).into());
}
if lhs_shape[1] != rhs_shape[0] {
return Err(
tenferro_tensor::Error::shape_mismatch("matmul", lhs_shape, rhs_shape).into(),
);
}
self.dot_general(
lhs,
rhs,
DotGeneralConfig {
lhs_contracting_dims: [1].as_slice().into(),
rhs_contracting_dims: [0].as_slice().into(),
lhs_batch_dims: [].as_slice().into(),
rhs_batch_dims: [].as_slice().into(),
},
)
}
pub fn dot_general_with_conj(
&mut self,
lhs: &EagerTensor,
rhs: &EagerTensor,
config: DotGeneralConfig,
lhs_conj: bool,
rhs_conj: bool,
) -> Result<EagerTensor> {
self.ensure_runtime(lhs)?;
self.ensure_runtime(rhs)?;
config
.validate_dims_with_ranks(lhs.shape().len(), rhs.shape().len())
.map_err(Error::TensorRuntime)?;
if !lhs.requires_grad && !rhs.requires_grad {
let output = crate::eager_exec::exec_dot_general_with_conj_on_tensor_reads_in_session(
lhs.tensor_read(),
rhs.tensor_read(),
&config,
lhs_conj,
rhs_conj,
&mut *self.engine,
)?;
return EagerTensor::new_untracked_result(Arc::clone(self.runtime), output);
}
let lhs = if lhs_conj {
self.conj(lhs)?
} else {
lhs.clone()
};
let rhs = if rhs_conj {
self.conj(rhs)?
} else {
rhs.clone()
};
self.dot_general(&lhs, &rhs, config)
}
fn run_binary(
&mut self,
name: &'static str,
lhs: &EagerTensor,
rhs: &EagerTensor,
op: StdTensorOp,
) -> Result<EagerTensor> {
self.ensure_runtime(lhs)?;
self.ensure_runtime(rhs)?;
let (lhs, rhs) = crate::eager_ops::broadcast_binary_in_session(name, lhs, rhs, self)?;
EagerTensor::nary_op_in_session(&[lhs.as_ref(), rhs.as_ref()], op, &mut *self.engine)
}
fn run_ternary(
&mut self,
name: &'static str,
first: &EagerTensor,
second: &EagerTensor,
third: &EagerTensor,
op: StdTensorOp,
) -> Result<EagerTensor> {
self.ensure_runtime(first)?;
self.ensure_runtime(second)?;
self.ensure_runtime(third)?;
let (first, second, third) =
crate::eager_ops::broadcast_ternary_in_session(name, first, second, third, self)?;
EagerTensor::nary_op_in_session(&[&first, &second, &third], op, &mut *self.engine)
}
pub fn apply_standard_op(
&mut self,
op: StdTensorOp,
inputs: &[&EagerTensor],
) -> Result<EagerTensor> {
if matches!(op, StdTensorOp::Extension(_)) {
return Err(Error::invalid_argument(
"EagerSession::apply_standard_op",
ErrorPhase::Execution,
"op",
"Extension ops must be passed to apply_eager",
));
}
for input in inputs {
self.ensure_runtime(input)?;
}
EagerTensor::nary_op_in_session(inputs, op, &mut *self.engine)
}
pub fn backend_session(&mut self) -> &mut dyn Engine {
&mut *self.engine
}
pub fn engine(&mut self) -> &mut dyn ExecutionEngine {
&mut *self.engine
}
fn run_unary(&mut self, input: &EagerTensor, op: StdTensorOp) -> Result<EagerTensor> {
self.ensure_runtime(input)?;
EagerTensor::nary_op_in_session(&[input], op, &mut *self.engine)
}
pub(crate) fn record_outputs(
&mut self,
op: &StdTensorOp,
outputs: &[&Tensor],
inputs: &[&EagerTensor],
) -> Result<RecordedEagerOutputs> {
record_eager_outputs_in_session(op, outputs, inputs, &mut *self.engine)
}
pub(crate) fn ensure_runtime(&self, input: &EagerTensor) -> Result<()> {
if !Arc::ptr_eq(self.runtime, &input.ctx) {
return Err(Error::ContextMismatch {
lhs: self.runtime.id(),
rhs: input.ctx_id(),
});
}
Ok(())
}
pub(crate) fn runtime(&self) -> &Arc<EagerRuntime> {
self.runtime
}
pub fn with_extension_caches<R>(
&mut self,
f: impl FnOnce(&mut tenferro_runtime::ExtensionCacheStore) -> R,
) -> Result<R> {
let result = with_lent_extension_caches(self.engine, |_, caches| f(caches));
Ok(result)
}
}
impl fmt::Debug for EagerRuntime {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
let mut debug = f.debug_struct("EagerRuntime");
debug.field("id", &self.id);
debug.field("runtime_id", &self.runtime.id());
debug.field("runtime_epoch", &self.runtime.epoch().ok());
match self.extension_install_lock.try_lock() {
Ok(_) => {
debug.field("extension_install_lock", &"<unlocked>");
}
Err(_) => {
debug.field("extension_install_lock", &"<locked>");
}
}
debug.field("semantic_extension_rules", &self.semantic_extension_rules);
match self.grad_slots.try_lock() {
Ok(slots) => {
debug.field("grad_slots_len", &slots.len());
}
Err(_) => {
debug.field("grad_slots_len", &"<locked>");
}
}
match self.value_records.try_lock() {
Ok(records) => {
debug.field("value_records_len", &records.len());
}
Err(_) => {
debug.field("value_records_len", &"<locked>");
}
}
match self.ad_transform_cache.stats() {
Ok(stats) => {
debug.field("ad_transform_cache_stats", &stats);
}
Err(err) => {
debug.field("ad_transform_cache_stats", &format_args!("{err}"));
}
}
match self.prepared_derivative_cache.try_lock() {
Ok(cache) => {
debug.field("prepared_derivative_cache_stats", &cache.stats());
}
Err(_) => {
debug.field("prepared_derivative_cache_stats", &"<locked>");
}
}
debug.finish_non_exhaustive()
}
}
impl EagerRuntime {
fn lock_extension_install(&self) -> Result<MutexGuard<'_, ()>> {
self.extension_install_lock.lock().map_err(|_| {
Error::runtime_state(
"eager_extension_install",
ErrorPhase::Execution,
"lock poisoned",
)
})
}
fn lock_prepared_derivative_cache(&self) -> Result<MutexGuard<'_, PreparedDerivativeCache>> {
self.prepared_derivative_cache.lock().map_err(|_| {
Error::runtime_state(
"prepared_derivative_cache",
ErrorPhase::Execution,
"lock poisoned",
)
})
}
fn lock_grad_slots(
&self,
) -> Result<MutexGuard<'_, HashMap<ValueKey<StdTensorOp>, WeakGradSlot>>> {
self.grad_slots.lock().map_err(|_| {
Error::runtime_state(
"eager_gradient_slots",
ErrorPhase::Execution,
"lock poisoned",
)
})
}
fn lock_value_records(
&self,
) -> Result<MutexGuard<'_, HashMap<ValueKey<StdTensorOp>, Weak<EagerTensorRecord>>>> {
self.value_records.lock().map_err(|_| {
Error::runtime_state(
"eager_value_registry",
ErrorPhase::Execution,
"lock poisoned",
)
})
}
fn with_rules_and_cache(
semantic_extension_rules: SemanticExtensionRuleSet,
ad_transform_cache: Arc<AdTransformCache>,
) -> Result<Self> {
let runtime = Runtime::builder()
.build()
.map_err(|source| runtime_config_error("EagerRuntime::new", source))?;
Ok(Self {
id: ContextId::fresh(),
runtime,
extension_install_lock: Mutex::new(()),
semantic_extension_rules,
grad_slots: Mutex::new(HashMap::new()),
value_records: Mutex::new(HashMap::new()),
ad_transform_cache,
prepared_derivative_cache: Mutex::new(PreparedDerivativeCache::default()),
})
}
pub fn new() -> Result<Arc<Self>> {
Ok(Arc::new(Self::with_rules_and_cache(
SemanticExtensionRuleSet::default(),
Arc::new(AdTransformCache::new()),
)?))
}
pub fn with_ad_context(ad: &AdContext) -> Result<Arc<Self>> {
Ok(Arc::new(Self::with_rules_and_cache(
ad.semantic_extension_rules().clone(),
ad.ad_transform_cache(),
)?))
}
pub(crate) fn ensure_engine(&self, engine: &mut dyn ExecutionEngine) -> Result<EngineTarget> {
let target = engine.target();
let snapshot = self
.runtime
.snapshot()
.map_err(|source| runtime_state_source("EagerRuntime::ensure_engine", source))?;
if let Some(registered) = snapshot.engine(target.engine_id()) {
if registered.context_identity() != target.context_identity() {
return Err(Error::unsupported(
"EagerRuntime::ensure_engine",
ErrorPhase::Execution,
format!(
"engine id {:?} is registered for a different engine type",
target.engine_id()
),
));
}
return Ok(target);
}
let registration = engine_kind_registration(engine)
.ok_or_else(|| unregisterable_engine_error(target.engine_id()))?;
let registration = registration
.map_err(|source| runtime_config_error("EagerRuntime::ensure_engine", source))?;
let _install_guard = self.lock_extension_install()?;
let snapshot = self
.runtime
.snapshot()
.map_err(|source| runtime_state_source("EagerRuntime::ensure_engine", source))?;
if snapshot.engine(target.engine_id()).is_none() {
self.runtime
.reconfigure(|edit| {
edit.register_engine(registration)?;
Ok(())
})
.map_err(|source| runtime_state_source("EagerRuntime::ensure_engine", source))?;
}
Ok(target)
}
pub fn id(&self) -> ContextId {
self.id
}
pub fn no_grad(&self) -> EagerNoGradGuard {
EAGER_NO_GRAD_DEPTH.with(|depth| {
depth.set(depth.get().saturating_add(1));
});
EagerNoGradGuard {
active: true,
_not_send: PhantomData,
}
}
pub fn capture_trace(&self) -> EagerTraceCaptureGuard {
EAGER_CAPTURE_DEPTH.with(|depth| {
depth.set(depth.get().saturating_add(1));
});
EagerTraceCaptureGuard {
active: true,
_not_send: PhantomData,
}
}
pub fn install_extension_module(
&self,
module: Arc<dyn ExtensionModule>,
) -> Result<RuntimeEpoch> {
let snapshot = self.runtime.snapshot().map_err(|source| {
runtime_state_source("EagerRuntime::install_extension_module", source)
})?;
if snapshot.has_extension_module_identical(&module) {
return Ok(snapshot.epoch());
}
let _install_guard = self.lock_extension_install()?;
self.runtime
.reconfigure(|edit| {
edit.replace_extension_module(module)?;
Ok(())
})
.map_err(|source| {
runtime_state_source("EagerRuntime::install_extension_module", source)
})
}
pub(crate) fn ensure_extension_module_for_engine(
&self,
module: Arc<dyn ExtensionModule>,
family_id: &'static str,
engine_id: &EngineId,
) -> Result<RuntimeEpoch> {
let snapshot = self.runtime.snapshot().map_err(|source| {
runtime_state_source("EagerRuntime::ensure_extension_module_for_engine", source)
})?;
if snapshot.has_extension_module_engine(module.module_id(), family_id, engine_id) {
return Ok(snapshot.epoch());
}
let _install_guard = self.lock_extension_install()?;
self.runtime
.reconfigure(|edit| {
edit.ensure_extension_module_for_engine(module, family_id, engine_id)?;
Ok(())
})
.map_err(|source| {
runtime_state_source("EagerRuntime::ensure_extension_module_for_engine", source)
})
}
pub(crate) fn runtime(&self) -> &Runtime {
&self.runtime
}
pub(crate) fn eager_extension_target(
&self,
engine: &mut dyn ExecutionEngine,
) -> Result<EagerExtensionTarget> {
let target = self.ensure_engine(engine)?;
let backend_kind = eager_extension_backend_kind(&target).ok_or_else(|| {
Error::unsupported(
"EagerRuntime::eager_extension_target",
ErrorPhase::Execution,
"this engine kind has no eager extension target",
)
})?;
let target = EagerExtensionTarget {
engine_id: target.engine_id().clone(),
backend_kind,
};
validate_eager_extension_target(&self.runtime, &target)?;
Ok(target)
}
pub fn clear_caches(&self) -> Result<()> {
self.clear_ad_transform_caches()?;
self.clear_prepared_derivative_cache()?;
Ok(())
}
pub fn clear_prepared_derivative_cache(&self) -> Result<()> {
self.lock_prepared_derivative_cache()?.clear();
Ok(())
}
pub fn cache_stats(&self) -> Result<EagerRuntimeCacheStats> {
Ok(EagerRuntimeCacheStats {
ad_transforms: self.ad_transform_cache.stats()?,
prepared_derivatives: self.lock_prepared_derivative_cache()?.stats(),
})
}
pub fn ad_transform_cache_limits(&self) -> Result<AdTransformCacheLimits> {
self.ad_transform_cache.limits()
}
pub fn set_ad_transform_cache_limits(&self, limits: AdTransformCacheLimits) -> Result<()> {
self.ad_transform_cache.set_limits(limits)
}
pub fn clear_ad_transform_caches(&self) -> Result<()> {
self.ad_transform_cache.clear()
}
pub fn prepared_derivative_cache_limits(&self) -> Result<AdTransformCacheLimits> {
Ok(self.lock_prepared_derivative_cache()?.limits())
}
pub fn set_prepared_derivative_cache_limits(
&self,
limits: AdTransformCacheLimits,
) -> Result<()> {
self.lock_prepared_derivative_cache()?.set_limits(limits);
Ok(())
}
pub fn with_eager_session<T, E: From<Error>>(
self: &Arc<Self>,
engine: &mut dyn ExecutionEngine,
f: impl FnOnce(&mut EagerSession<'_>) -> std::result::Result<T, E>,
) -> std::result::Result<T, E> {
self.ensure_engine(engine).map_err(E::from)?;
f(&mut EagerSession {
runtime: self,
engine,
})
}
pub(crate) fn run_program(
&self,
engine: &mut dyn ExecutionEngine,
program: &CompiledGraph,
inputs: &[TensorRead<'_>],
) -> Result<Vec<Tensor>> {
let target = self.ensure_engine(engine)?;
let prepared = self
.runtime
.prepare(program, inputs, std::slice::from_ref(&target))?;
run_prepared_on(&self.runtime, engine, &prepared, inputs)
}
pub(crate) fn exec_extension_outputs_read(
&self,
engine: &mut dyn ExecutionEngine,
op: &Arc<dyn tenferro_ops::ext_op::ExtensionOp>,
inputs: &[TensorRead<'_>],
) -> Result<Vec<Tensor>> {
self.ensure_engine(engine)?;
profile_eager_op_section("exec_extension_outputs_read.exec_op", || {
exec_extension_op_on_tensor_reads(op, inputs, engine, &self.runtime)
})
}
#[cfg(test)]
pub(crate) fn exec_standard_graph_outputs(
&self,
engine: &mut dyn ExecutionEngine,
graph: &Graph<StdTensorOp>,
initial_data: HashMap<ValueKey<StdTensorOp>, Tensor>,
) -> Result<EagerGraphExecution> {
let mut all_values = initial_data;
profile_eager_op_section("exec_graph.run", || {
let exec: &mut dyn Engine = engine;
(|| -> Result<()> {
for op_node in graph.operations() {
let outputs = {
let input_values = op_node
.inputs
.iter()
.map(|input| {
let key = match input {
ValueRef::Local(local_id) => &graph.values()[*local_id].key,
ValueRef::External(key) => key,
};
all_values.get(key).ok_or_else(|| {
Error::Internal(format!(
"standard graph eager execution missing value for {key:?}"
))
})
})
.collect::<Result<Vec<_>>>()?;
let input_reads = input_values
.iter()
.map(|value| TensorRead::from_tensor(value))
.collect::<Vec<_>>();
exec_standard_op_on_tensor_reads_in_session(
&op_node.operation,
&input_reads,
exec,
)?
};
if outputs.len() != op_node.outputs.len() {
return Err(Error::Internal(format!(
"standard graph eager execution expected {} outputs for {:?}, got {}",
op_node.outputs.len(),
op_node.operation,
outputs.len()
)));
}
for (output_id, output) in op_node.outputs.iter().zip(outputs) {
let key = graph.values()[*output_id].key.clone();
all_values.insert(key, output);
}
}
Ok(())
})()
})?;
let outputs = graph
.outputs()
.iter()
.map(|&output_id| {
let key = &graph.values()[output_id].key;
all_values
.get(key)
.ok_or_else(|| {
Error::Internal(format!(
"standard graph eager execution missing graph output {key:?}"
))
})?
.duplicate()
.map_err(Error::from)
})
.collect::<Result<Vec<_>>>()?;
Ok(EagerGraphExecution { outputs })
}
pub(crate) fn try_register_grad_slot(
&self,
key: &ValueKey<StdTensorOp>,
slot: &GradSlot,
) -> Result<()> {
insert_pruning_dead(
&mut *self.lock_grad_slots()?,
key.clone(),
Arc::downgrade(slot),
);
Ok(())
}
pub(crate) fn try_register_value_record(
&self,
key: &ValueKey<StdTensorOp>,
record: &Arc<EagerTensorRecord>,
) -> Result<()> {
insert_pruning_dead(
&mut *self.lock_value_records()?,
key.clone(),
Arc::downgrade(record),
);
Ok(())
}
pub(crate) fn value_record(
&self,
key: &ValueKey<StdTensorOp>,
) -> Result<Option<Arc<EagerTensorRecord>>> {
let mut records = self.lock_value_records()?;
let Some(record) = records.get(key).cloned() else {
return Ok(None);
};
match record.upgrade() {
Some(record) => Ok(Some(record)),
None => {
records.remove(key);
Ok(None)
}
}
}
pub fn clear_grads(&self) -> Result<()> {
let live_slots = {
let mut live_slots = Vec::new();
self.lock_grad_slots()?.retain(|_, slot| {
if let Some(slot) = slot.upgrade() {
live_slots.push(slot);
true
} else {
false
}
});
live_slots
};
let mut poisoned_slot = false;
for slot in live_slots {
match slot.lock() {
Ok(mut current) => {
*current = None;
}
Err(_) => {
poisoned_slot = true;
}
}
}
if poisoned_slot {
return Err(Error::runtime_state(
"eager_gradient_slot",
ErrorPhase::Execution,
"lock poisoned",
));
}
Ok(())
}
pub fn constant_from(self: &Arc<Self>, tensor: Tensor) -> Result<EagerTensor> {
EagerTensor::new_leaf(Arc::clone(self), tensor, false)
}
pub fn variable_from(self: &Arc<Self>, tensor: Tensor) -> Result<EagerTensor> {
EagerTensor::new_leaf(Arc::clone(self), tensor, true)
}
pub fn grad(
self: &Arc<Self>,
engine: &mut dyn ExecutionEngine,
output: &EagerTensor,
wrt: &EagerTensor,
) -> Result<EagerTensor> {
self.grad_optional(engine, output, wrt)?
.ok_or_else(|| crate::traced::inactive_wrt_error("grad", &wrt.key))
}
pub fn grad_optional(
self: &Arc<Self>,
engine: &mut dyn ExecutionEngine,
output: &EagerTensor,
wrt: &EagerTensor,
) -> Result<Option<EagerTensor>> {
if !output.shape().is_empty() {
return Err(Error::NonScalarGrad {
shape: output.shape().to_vec(),
});
}
let value = output.to_tensor(&mut *engine)?;
let seed = one_like_tensor(&value, &mut *engine)?;
let seed = EagerTensor::new_result(Arc::clone(self), eager_val_key(), seed, false, None)?;
self.vjp_optional(engine, output, wrt, &seed)
}
pub fn vjp(
self: &Arc<Self>,
engine: &mut dyn ExecutionEngine,
output: &EagerTensor,
wrt: &EagerTensor,
cotangent: &EagerTensor,
) -> Result<EagerTensor> {
self.vjp_optional(engine, output, wrt, cotangent)?
.ok_or_else(|| crate::traced::inactive_wrt_error("vjp", &wrt.key))
}
pub fn vjp_optional(
self: &Arc<Self>,
engine: &mut dyn ExecutionEngine,
output: &EagerTensor,
wrt: &EagerTensor,
cotangent: &EagerTensor,
) -> Result<Option<EagerTensor>> {
validate_same_runtime(self, output, "vjp output")?;
validate_same_runtime(self, wrt, "vjp wrt")?;
validate_same_runtime(self, cotangent, "vjp cotangent")?;
validate_seed_tensor("vjp", output, cotangent)?;
Ok(
semantic_eager_vjp_many(self, engine, output, &[wrt], cotangent)?
.pop()
.flatten(),
)
}
pub fn jvp(
self: &Arc<Self>,
engine: &mut dyn ExecutionEngine,
output: &EagerTensor,
wrt: &EagerTensor,
tangent: &EagerTensor,
) -> Result<EagerTensor> {
self.jvp_optional(engine, output, wrt, tangent)?
.ok_or_else(|| crate::traced::inactive_wrt_error("jvp", &wrt.key))
}
pub fn jvp_optional(
self: &Arc<Self>,
engine: &mut dyn ExecutionEngine,
output: &EagerTensor,
wrt: &EagerTensor,
tangent: &EagerTensor,
) -> Result<Option<EagerTensor>> {
validate_same_runtime(self, output, "jvp output")?;
validate_same_runtime(self, wrt, "jvp wrt")?;
validate_same_runtime(self, tangent, "jvp tangent")?;
validate_seed_tensor("jvp", wrt, tangent)?;
match semantic_eager_jvp_optional(self, engine, output, wrt, tangent)? {
Some(result) => Ok(result),
None => Ok(None),
}
}
fn store_grads(
&self,
cotangents: &HashMap<ValueKey<StdTensorOp>, Tensor>,
session: &mut dyn Engine,
) -> Result<()> {
let mut updates = Vec::new();
{
let mut slots = self.lock_grad_slots()?;
slots.retain(|key, slot| {
let Some(slot) = slot.upgrade() else {
return false;
};
if let Some(incoming) = cotangents.get(key) {
updates.push((slot, incoming));
}
true
});
}
for (slot, incoming) in updates {
let mut current = slot.lock().map_err(|_| {
Error::runtime_state(
"eager_gradient_slot",
ErrorPhase::Execution,
"lock poisoned",
)
})?;
let next = match current.as_ref() {
Some(existing) => {
let existing_read = existing.tensor_read("EagerRuntime::store_grads")?;
let incoming_read = TensorRead::from_tensor(incoming);
let tensor = session
.add_read(existing_read, incoming_read)
.map_err(Error::from)?;
AdValueRecord::from_tensor(tensor, "EagerRuntime::store_grads")?
}
None => {
let duplicate = session
.to_contiguous_read(TensorRead::from_tensor(incoming))
.map_err(Error::from)?;
AdValueRecord::from_tensor(duplicate, "EagerRuntime::store_grads")?
}
};
*current = Some(next);
}
Ok(())
}
}
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
struct PreparedDerivativeCacheKey {
target: EngineTarget,
semantic_fingerprint: SemanticFingerprint,
runtime_epoch: RuntimeEpoch,
active_inputs: Box<[bool]>,
input_metadata: Box<[ProgramValueMetadata]>,
}
#[derive(Debug)]
struct PreparedDerivative {
program: Arc<CompiledGraph>,
execution_program: Arc<CompiledGraph>,
saved_input_indices: Vec<usize>,
prepared: Arc<Prepared>,
seed_input_index: usize,
derivative_output_indices: Box<[Option<usize>]>,
}
#[derive(Debug)]
struct PreparedDerivativeCache {
limits: AdTransformCacheLimits,
entries: LruCache<PreparedDerivativeCacheKey, PreparedDerivativeCacheEntry>,
stats: CacheStats,
}
impl PreparedDerivativeCache {
fn limits(&self) -> AdTransformCacheLimits {
self.limits
}
fn set_limits(&mut self, limits: AdTransformCacheLimits) {
self.limits = limits;
self.evict_to_limits();
}
fn clear(&mut self) {
let clears = self.stats.clears.saturating_add(1);
self.entries.clear();
self.stats = CacheStats {
clears,
..CacheStats::empty()
};
}
fn stats(&self) -> CacheStats {
self.stats
}
fn get(&mut self, key: &PreparedDerivativeCacheKey) -> Option<Arc<PreparedDerivative>> {
match self.entries.get(key) {
Some(entry) => {
self.stats.hits = self.stats.hits.saturating_add(1);
Some(Arc::clone(&entry.value))
}
None => {
self.stats.misses = self.stats.misses.saturating_add(1);
None
}
}
}
fn insert(&mut self, key: PreparedDerivativeCacheKey, value: Arc<PreparedDerivative>) {
let retained_bytes = prepared_derivative_cache_entry_retained_bytes(&key, value.as_ref());
let entry = PreparedDerivativeCacheEntry {
value,
retained_bytes,
};
self.stats.retained_bytes = self.stats.retained_bytes.saturating_add(retained_bytes);
if let Some((_old_key, old_entry)) = self.entries.push(key, entry) {
self.stats.retained_bytes = self
.stats
.retained_bytes
.saturating_sub(old_entry.retained_bytes);
}
self.stats.entries = self.entries.len();
self.evict_to_limits();
}
fn evict_to_limits(&mut self) {
while self.entries.len() > self.limits.max_entries().get()
|| self
.limits
.max_retained_bytes()
.is_some_and(|limit| self.stats.retained_bytes > limit.get())
{
let Some((_key, entry)) = self.entries.pop_lru() else {
break;
};
self.stats.retained_bytes = self
.stats
.retained_bytes
.saturating_sub(entry.retained_bytes);
self.stats.evictions = self.stats.evictions.saturating_add(1);
}
self.stats.entries = self.entries.len();
}
}
impl Default for PreparedDerivativeCache {
fn default() -> Self {
Self {
limits: AdTransformCacheLimits::default(),
entries: LruCache::unbounded(),
stats: CacheStats::empty(),
}
}
}
#[derive(Debug)]
struct PreparedDerivativeCacheEntry {
value: Arc<PreparedDerivative>,
retained_bytes: usize,
}
fn prepared_derivative_cache_entry_retained_bytes(
key: &PreparedDerivativeCacheKey,
value: &PreparedDerivative,
) -> usize {
size_of::<PreparedDerivativeCacheKey>()
.saturating_add(size_of_val(key.active_inputs.as_ref()))
.saturating_add(size_of_val(value.derivative_output_indices.as_ref()))
.saturating_add(size_of_val(value.saved_input_indices.as_slice()))
.saturating_add(compiled_graph_retained_bytes(
value.execution_program.as_ref(),
))
.saturating_add(
key.input_metadata
.len()
.saturating_mul(size_of::<ProgramValueMetadata>()),
)
.saturating_add(size_of::<PreparedDerivative>())
.saturating_add(compiled_graph_retained_bytes(value.program.as_ref()))
.saturating_add(prepared_compiled_graph_retained_bytes(
value.prepared.as_ref(),
value.program.as_ref(),
))
}
fn prepared_compiled_graph_retained_bytes(
prepared: &Prepared,
derivative_program: &CompiledGraph,
) -> usize {
size_of_val(prepared).saturating_add(compiled_graph_retained_bytes(derivative_program))
}
fn compiled_graph_retained_bytes(program: &CompiledGraph) -> usize {
size_of::<CompiledGraph>()
.saturating_add(size_of_val(program.input_keys()))
.saturating_add(program.bindings().len().saturating_mul(size_of::<usize>()))
.saturating_add(semantic_program_retained_bytes(program.program()))
}
fn semantic_program_retained_bytes(program: &SemanticProgram) -> usize {
size_of::<SemanticProgram>()
.saturating_add(size_of_val(program.inputs()))
.saturating_add(size_of_val(program.outputs()))
.saturating_add(
program
.operations()
.len()
.saturating_mul(size_of::<usize>()),
)
.saturating_add(
program
.shape_guards()
.len()
.saturating_mul(size_of::<usize>()),
)
}
fn semantic_eager_vjp_many(
ctx: &Arc<EagerRuntime>,
engine: &mut dyn ExecutionEngine,
output: &EagerTensor,
wrts: &[&EagerTensor],
cotangent: &EagerTensor,
) -> Result<Vec<Option<EagerTensor>>> {
if !eager_semantic_vjp_enabled() || wrts.is_empty() {
return Ok(vec![None; wrts.len()]);
}
let Some(raw_output_trace) = output.semantic_trace.as_ref() else {
return Ok(vec![None; wrts.len()]);
};
if !wrts.iter().any(|wrt| {
wrt.semantic_trace
.as_ref()
.and_then(TracedTensor::input_key)
.is_some_and(|key| raw_output_trace.has_attached_input_key(&key))
}) {
return Ok(vec![None; wrts.len()]);
}
let output_trace = analyze_deferred_semantic_trace(raw_output_trace)?;
let saved = output
.trace
.as_ref()
.map(EagerTrace::collect)
.unwrap_or_default();
let mut source_outputs = vec![&output_trace];
source_outputs.extend(saved.iter().map(|value| &value.trace));
let mut compiler = GraphCompiler::new();
let source =
tenferro_runtime::ad_support::compile_ad_source_many(&mut compiler, &source_outputs)?;
if source.output_count() != source_outputs.len()
|| source.input_keys().len() != source.input_count()
|| source.bindings().len() != source.input_count()
{
return Ok(vec![None; wrts.len()]);
}
let wrt_input_indices = wrts
.iter()
.map(|wrt| {
let key = wrt.semantic_trace.as_ref()?.input_key()?;
source.input_key_index(&key)
})
.collect::<Vec<_>>();
let mut active_inputs = vec![false; source.input_count()];
for &index in wrt_input_indices.iter().flatten() {
active_inputs[index] = true;
}
if !active_inputs.iter().any(|&active| active) {
return Ok(vec![None; wrts.len()]);
}
let target = ctx.ensure_engine(engine)?;
let cache_key = PreparedDerivativeCacheKey {
target: target.clone(),
semantic_fingerprint: source.program().semantic_fingerprint(),
runtime_epoch: ctx.runtime.epoch().map_err(|source| {
Error::runtime_state_source("semantic_eager_vjp", ErrorPhase::Execution, source)
})?,
active_inputs: active_inputs.clone().into_boxed_slice(),
input_metadata: source.frozen_program().input_metadata_with_bound_shapes(),
};
let prepared = { ctx.lock_prepared_derivative_cache()?.get(&cache_key) };
let (
seed_input_index,
derivative_output_indices,
derivative_program,
execution_program,
saved_input_indices,
prepared_runtime,
) = if let Some(prepared) = prepared {
(
prepared.seed_input_index,
prepared.derivative_output_indices.clone(),
Arc::clone(&prepared.program),
Arc::clone(&prepared.execution_program),
prepared.saved_input_indices.clone(),
Some(Arc::clone(&prepared.prepared)),
)
} else {
let mut active_outputs = vec![false; source.output_count()];
active_outputs[0] = true;
let ad = AdContext::with_rules_and_transform_cache(
ctx.semantic_extension_rules.clone(),
Arc::clone(&ctx.ad_transform_cache),
);
let derivative = ad
.vjp_program(source.frozen_program(), &active_inputs, &active_outputs)
.map_err(|source| {
Error::runtime_state_source("semantic_eager_vjp", ErrorPhase::GraphBuild, source)
})?;
let seed_input_index = derivative
.derivative_input_indices()
.first()
.copied()
.flatten();
let Some(seed_input_index) = seed_input_index else {
return Ok(vec![None; wrts.len()]);
};
let derivative_output_indices = derivative
.derivative_output_indices()
.to_vec()
.into_boxed_slice();
let program = Arc::new(compiler.compile_frozen_program(derivative.frozen())?);
let (execution_program, saved_input_indices) = if saved.is_empty() {
(Arc::clone(&program), Vec::new())
} else {
let (execution, indices) = crate::semantic_transform::semantic_vjp_with_saved_outputs(
source.frozen_program(),
&active_inputs,
&active_outputs,
&ctx.semantic_extension_rules,
&(1..source.output_count()).collect::<Vec<_>>(),
)
.map_err(|source| {
Error::runtime_state_source("semantic_eager_vjp", ErrorPhase::GraphBuild, source)
})?;
(
Arc::new(compiler.compile_frozen_program(execution.frozen())?),
indices,
)
};
(
seed_input_index,
derivative_output_indices,
program,
execution_program,
saved_input_indices,
None,
)
};
let cotangent_tensor = Arc::new(RetainedValue::from_tensor(
cotangent.to_tensor(&mut *engine)?,
));
let input_count = execution_program.input_count();
let mut owned_inputs: Vec<Option<Tensor>> = (0..input_count).map(|_| None).collect();
(|| -> Result<()> {
let session: &mut dyn Engine = &mut *engine;
for (value, &index) in saved.iter().zip(&saved_input_indices) {
let read = value.value.tensor_read("eager residual")?;
let Some(slot) = owned_inputs.get_mut(index) else {
return Err(Error::Internal(format!(
"semantic eager VJP residual index {index} is outside {input_count} inputs"
)));
};
*slot = Some(session.to_contiguous_read(read)?);
}
for (source_input_index, (_, tensor)) in source.bindings().iter().enumerate() {
let Some(slot) = owned_inputs.get_mut(source_input_index) else {
return Err(Error::Internal(format!(
"semantic eager VJP derivative program has no primal input slot {source_input_index}"
)));
};
*slot = Some(copy_value_in_session(session, tensor)?);
}
let Some(slot) = owned_inputs.get_mut(seed_input_index) else {
return Err(Error::Internal(format!(
"semantic eager VJP seed input index {seed_input_index} is outside {input_count} inputs"
)));
};
*slot = Some(copy_value_in_session(session, cotangent_tensor.as_ref())?);
Ok(())
})()?;
let input_refs = owned_inputs
.iter()
.enumerate()
.map(|(index, tensor)| {
tensor.as_ref().ok_or_else(|| {
Error::Internal(format!(
"semantic eager VJP derivative input {index} was not populated"
))
})
})
.collect::<Result<Vec<_>>>()?;
let input_reads = input_refs
.iter()
.map(|tensor| TensorRead::from_tensor(tensor))
.collect::<Vec<_>>();
let prepared_runtime = if let Some(prepared_runtime) = prepared_runtime {
prepared_runtime
} else {
let prepared_runtime = Arc::new(ctx.runtime.prepare(
&execution_program,
&input_reads,
std::slice::from_ref(&target),
)?);
let entry = Arc::new(PreparedDerivative {
program: Arc::clone(&derivative_program),
execution_program: Arc::clone(&execution_program),
saved_input_indices: saved_input_indices.clone(),
prepared: Arc::clone(&prepared_runtime),
seed_input_index,
derivative_output_indices: derivative_output_indices.clone(),
});
ctx.lock_prepared_derivative_cache()?
.insert(cache_key, entry);
prepared_runtime
};
let mut outputs = run_prepared_on(&ctx.runtime, engine, &prepared_runtime, &input_reads)?
.into_iter()
.map(Some)
.collect::<Vec<_>>();
let cotangent_trace =
TracedTensor::from_shared_tensor_value_symbolic_shape(Arc::clone(&cotangent_tensor))?;
#[cfg(test)]
EAGER_SEMANTIC_VJP_EXECUTIONS.fetch_add(1, Ordering::Relaxed);
wrts.iter()
.zip(wrt_input_indices)
.map(|(wrt, input_index)| {
let Some(derivative_output_index) = input_index
.and_then(|index| derivative_output_indices.get(index).copied().flatten())
else {
return Ok(None);
};
let result = outputs
.get_mut(derivative_output_index)
.and_then(Option::take)
.ok_or_else(|| {
Error::Internal(format!(
"semantic eager VJP derivative output {derivative_output_index} unavailable"
))
})?;
let wrt_trace = wrt.semantic_trace.as_ref().ok_or_else(|| {
Error::Internal("active eager VJP input has no semantic trace".into())
})?;
let semantic_trace = derivative_trace_from_frozen_program(
&source,
derivative_program.frozen_program(),
derivative_output_index,
&[(seed_input_index, Arc::clone(&cotangent_tensor))],
&[&output_trace, wrt_trace, &cotangent_trace],
None,
"semantic_eager_vjp",
)?;
Ok(Some(EagerTensor::new_result_with_semantic_trace(
Arc::clone(ctx),
eager_val_key(),
result,
true,
output.trace.clone(),
Some(semantic_trace),
)?))
})
.collect()
}
fn semantic_eager_jvp_optional(
ctx: &Arc<EagerRuntime>,
engine: &mut dyn ExecutionEngine,
output: &EagerTensor,
wrt: &EagerTensor,
tangent: &EagerTensor,
) -> Result<Option<Option<EagerTensor>>> {
if !eager_semantic_vjp_enabled() {
return Ok(None);
}
let (Some(raw_output_trace), Some(wrt_trace)) =
(output.semantic_trace.as_ref(), wrt.semantic_trace.as_ref())
else {
return Ok(None);
};
let Some(wrt_key) = wrt_trace.input_key() else {
return Ok(None);
};
if !raw_output_trace.has_attached_input_key(&wrt_key) {
return Ok(None);
}
let output_trace = analyze_deferred_semantic_trace(raw_output_trace)?;
let mut compiler = GraphCompiler::new();
let source = compile_ad_source(&mut compiler, &output_trace)?;
if source.output_count() != 1
|| source.input_keys().len() != source.input_count()
|| source.bindings().len() != source.input_count()
{
return Ok(None);
}
let Some(wrt_input_index) = source.input_key_index(&wrt_key) else {
return Ok(None);
};
let mut active_inputs = vec![false; source.input_count()];
if let Some(active) = active_inputs.get_mut(wrt_input_index) {
*active = true;
} else {
return Ok(None);
}
let ad = AdContext::with_rules_and_transform_cache(
ctx.semantic_extension_rules.clone(),
Arc::clone(&ctx.ad_transform_cache),
);
let derivative = ad
.jvp_program(source.frozen_program(), &active_inputs)
.map_err(|source| {
Error::runtime_state_source("semantic_eager_jvp", ErrorPhase::GraphBuild, source)
})?;
let Some(seed_input_index) = derivative
.derivative_input_indices()
.get(wrt_input_index)
.copied()
.flatten()
else {
return Ok(Some(None));
};
let Some(derivative_output_index) = derivative
.derivative_output_indices()
.first()
.copied()
.flatten()
else {
return Ok(Some(None));
};
let derivative_program = compiler.compile_frozen_program(derivative.frozen())?;
let tangent_tensor = Arc::new(RetainedValue::from_tensor(tangent.to_tensor(&mut *engine)?));
let input_count = derivative_program.input_count();
let mut owned_inputs: Vec<Option<Tensor>> = (0..input_count).map(|_| None).collect();
(|| -> Result<()> {
let session: &mut dyn Engine = &mut *engine;
for (source_input_index, (_, tensor)) in source.bindings().iter().enumerate() {
let Some(slot) = owned_inputs.get_mut(source_input_index) else {
return Err(Error::Internal(format!(
"semantic eager JVP derivative program has no primal input slot {source_input_index}"
)));
};
*slot = Some(copy_value_in_session(session, tensor)?);
}
let Some(slot) = owned_inputs.get_mut(seed_input_index) else {
return Err(Error::Internal(format!(
"semantic eager JVP seed input index {seed_input_index} is outside {input_count} inputs"
)));
};
*slot = Some(copy_value_in_session(session, tangent_tensor.as_ref())?);
Ok(())
})()?;
let input_refs = owned_inputs
.iter()
.enumerate()
.map(|(index, tensor)| {
tensor.as_ref().ok_or_else(|| {
Error::Internal(format!(
"semantic eager JVP derivative input {index} was not populated"
))
})
})
.collect::<Result<Vec<_>>>()?;
let input_reads = input_refs
.iter()
.map(|tensor| TensorRead::from_tensor(tensor))
.collect::<Vec<_>>();
let outputs = ctx.run_program(engine, &derivative_program, &input_reads)?;
let output_count = outputs.len();
let Some(result) = outputs.into_iter().nth(derivative_output_index) else {
return Err(Error::Internal(format!(
"semantic eager JVP derivative output index {derivative_output_index} is outside {} outputs",
output_count
)));
};
let tangent_trace =
TracedTensor::from_shared_tensor_value_symbolic_shape(Arc::clone(&tangent_tensor))?;
let semantic_trace = derivative_trace_from_frozen_program(
&source,
derivative.frozen(),
derivative_output_index,
&[(seed_input_index, Arc::clone(&tangent_tensor))],
&[&output_trace, wrt_trace, &tangent_trace],
None,
"semantic_eager_jvp",
)?;
Ok(Some(Some(EagerTensor::new_result_with_semantic_trace(
Arc::clone(ctx),
eager_val_key(),
result,
true,
None,
Some(semantic_trace),
)?)))
}
fn validate_same_runtime(
runtime: &Arc<EagerRuntime>,
tensor: &EagerTensor,
role: &'static str,
) -> Result<()> {
if tensor.ctx_id() != runtime.id() {
return Err(Error::ContextMismatch {
lhs: runtime.id(),
rhs: tensor.ctx_id(),
});
}
let _ = role;
Ok(())
}
fn copy_value_in_session(session: &mut dyn Engine, value: &RetainedValue) -> Result<Tensor> {
let read = value.tensor_read().map_err(|error| {
Error::runtime_state_source("copy_value_for_runtime", ErrorPhase::Execution, error)
})?;
session.to_contiguous_read(read).map_err(Error::from)
}
fn validate_seed_tensor(op: &'static str, primal: &EagerTensor, seed: &EagerTensor) -> Result<()> {
if primal.dtype() != seed.dtype() {
return Err(
tenferro_tensor::Error::dtype_mismatch(op, primal.dtype(), seed.dtype()).into(),
);
}
if primal.shape() != seed.shape() {
return Err(
tenferro_tensor::Error::shape_mismatch(op, primal.shape(), seed.shape()).into(),
);
}
Ok(())
}
#[derive(Clone)]
pub struct EagerTensor {
pub(crate) key: ValueKey<StdTensorOp>,
pub(crate) trace: Option<EagerTrace>,
pub(crate) semantic_trace: Option<TracedTensor>,
pub(crate) requires_grad: bool,
grad_slot: GradSlot,
pub(crate) ctx: Arc<EagerRuntime>,
_record: Arc<EagerTensorRecord>,
}
pub(crate) struct EagerTensorRecord {
value: Arc<AdValueRecord>,
key: ValueKey<StdTensorOp>,
trace: Option<EagerTrace>,
semantic_trace: Option<TracedTensor>,
requires_grad: bool,
grad_slot: GradSlot,
ctx: Arc<EagerRuntime>,
}
struct EagerTensorParts {
ctx: Arc<EagerRuntime>,
key: ValueKey<StdTensorOp>,
requires_grad: bool,
trace: Option<EagerTrace>,
semantic_trace: Option<TracedTensor>,
value: Arc<AdValueRecord>,
register_value: bool,
}
impl fmt::Debug for EagerTensor {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("EagerTensor")
.field("dtype", &self.dtype())
.field("shape", &self.shape())
.field("key", &self.key)
.field("requires_grad", &self.requires_grad)
.field("has_trace", &self.trace.is_some())
.field("has_semantic_trace", &self.semantic_trace.is_some())
.field("ctx_id", &self.ctx_id())
.finish_non_exhaustive()
}
}
impl EagerTensor {
pub fn from_tensor_in(tensor: Tensor, ctx: Arc<EagerRuntime>) -> Result<Self> {
Self::new_leaf(ctx, tensor, false)
}
pub fn from_vec_col_major_in<T: TensorScalar>(
shape: impl IntoShapeVec,
data: Vec<T>,
ctx: Arc<EagerRuntime>,
) -> Result<Self> {
Self::from_tensor_in(Tensor::from_vec_col_major(shape, data)?, ctx)
}
pub fn requires_grad_in(tensor: Tensor, ctx: Arc<EagerRuntime>) -> Result<Self> {
Self::new_leaf(ctx, tensor, true)
}
pub(crate) fn new_leaf(
ctx: Arc<EagerRuntime>,
tensor: Tensor,
requires_grad: bool,
) -> Result<Self> {
Self::new_leaf_with_session(ctx, tensor, requires_grad, None)
}
pub(crate) fn new_leaf_in_session(
ctx: Arc<EagerRuntime>,
tensor: Tensor,
requires_grad: bool,
session: &mut dyn Engine,
) -> Result<Self> {
Self::new_leaf_with_session(ctx, tensor, requires_grad, Some(session))
}
fn new_leaf_with_session(
ctx: Arc<EagerRuntime>,
tensor: Tensor,
requires_grad: bool,
session: Option<&mut dyn Engine>,
) -> Result<Self> {
let key = eager_val_key();
let read = TensorRead::from_tensor(&tensor);
let semantic_tensor =
match session {
Some(session) => session.to_contiguous_read(read).map_err(Error::from)?,
None => match cpu_host_owned_read(&read) {
Some(materialized) => materialized.map_err(Error::from)?,
None => return Err(Error::unsupported(
"EagerTensor::new_leaf",
ErrorPhase::Execution,
"this tensor needs an engine to become an eager leaf; create it through \
EagerSession::constant_from or EagerSession::variable_from",
)),
},
};
let semantic_value = Arc::new(RetainedValue::from_tensor(semantic_tensor));
let semantic_trace = TracedTensor::from_shared_tensor_value_symbolic_shape(semantic_value)?;
let value = AdValueRecord::from_tensor(tensor, "EagerTensor::new_leaf")?;
Self::from_parts(EagerTensorParts {
ctx,
key,
requires_grad,
trace: None,
semantic_trace: Some(semantic_trace),
value,
register_value: true,
})
}
pub(crate) fn new_result(
ctx: Arc<EagerRuntime>,
key: ValueKey<StdTensorOp>,
tensor: Tensor,
requires_grad: bool,
trace: Option<EagerTrace>,
) -> Result<Self> {
Self::new_result_with_semantic_trace(ctx, key, tensor, requires_grad, trace, None)
}
pub(crate) fn new_result_with_semantic_trace(
ctx: Arc<EagerRuntime>,
key: ValueKey<StdTensorOp>,
tensor: Tensor,
requires_grad: bool,
trace: Option<EagerTrace>,
semantic_trace: Option<TracedTensor>,
) -> Result<Self> {
let value = AdValueRecord::from_tensor(tensor, "EagerTensor::new_result")?;
Self::from_parts(EagerTensorParts {
ctx,
key,
requires_grad,
trace,
semantic_trace,
value,
register_value: true,
})
}
pub(crate) fn new_unregistered_result_with_semantic_trace(
ctx: Arc<EagerRuntime>,
key: ValueKey<StdTensorOp>,
tensor: Tensor,
requires_grad: bool,
trace: Option<EagerTrace>,
semantic_trace: Option<TracedTensor>,
) -> Result<Self> {
let value = AdValueRecord::from_tensor(tensor, "EagerTensor::new_unregistered_result")?;
Self::from_parts(EagerTensorParts {
ctx,
key,
requires_grad,
trace,
semantic_trace,
value,
register_value: false,
})
}
pub(crate) fn new_result_value(
ctx: Arc<EagerRuntime>,
key: ValueKey<StdTensorOp>,
value: TensorValue,
requires_grad: bool,
trace: Option<EagerTrace>,
semantic_trace: Option<TracedTensor>,
) -> Result<Self> {
let (group, slot, dtype, shape) = value.try_into_group_parts().map_err(|_| {
Error::runtime_state(
"EagerTensor::new_result_value",
ErrorPhase::Execution,
"a TensorValue could not be transferred into its allocation group",
)
})?;
let value = AdValueRecord::from_group(group, slot, dtype, shape);
Self::from_parts(EagerTensorParts {
ctx,
key,
requires_grad,
trace,
semantic_trace,
value,
register_value: true,
})
}
fn from_parts(parts: EagerTensorParts) -> Result<Self> {
let EagerTensorParts {
ctx,
key,
requires_grad,
trace,
semantic_trace,
value,
register_value,
} = parts;
let grad_slot = Arc::new(Mutex::new(None));
if requires_grad {
ctx.try_register_grad_slot(&key, &grad_slot)?;
}
let record = Arc::new(EagerTensorRecord {
value: Arc::clone(&value),
key: key.clone(),
trace: trace.clone(),
semantic_trace: semantic_trace.clone(),
requires_grad,
grad_slot: Arc::clone(&grad_slot),
ctx: Arc::clone(&ctx),
});
if register_value {
ctx.try_register_value_record(&key, &record)?;
}
Ok(Self {
key,
trace,
semantic_trace,
requires_grad,
grad_slot,
ctx,
_record: record,
})
}
pub(crate) fn new_untracked_result(ctx: Arc<EagerRuntime>, tensor: Tensor) -> Result<Self> {
let value =
AdValueRecord::from_untracked_tensor(tensor, "EagerTensor::new_untracked_result")?;
Ok(Self::new_untracked_value_record(ctx, value, None))
}
pub(crate) fn new_untracked_value_result(
ctx: Arc<EagerRuntime>,
value: TensorValue,
) -> Result<Self> {
Self::new_untracked_value_result_with_semantic_trace(ctx, value, None)
}
pub(crate) fn new_untracked_value_result_with_semantic_trace(
ctx: Arc<EagerRuntime>,
value: TensorValue,
semantic_trace: Option<TracedTensor>,
) -> Result<Self> {
let (group, slot, dtype, shape) = value.try_into_group_parts().map_err(|_| {
Error::runtime_state(
"EagerTensor::new_untracked_value_result",
ErrorPhase::Execution,
"a TensorValue could not be transferred into its allocation group",
)
})?;
let value = AdValueRecord::from_group(group, slot, dtype, shape);
Ok(Self::new_untracked_value_record(ctx, value, semantic_trace))
}
fn new_untracked_value_record(
ctx: Arc<EagerRuntime>,
value: Arc<AdValueRecord>,
semantic_trace: Option<TracedTensor>,
) -> Self {
let key = eager_val_key();
let grad_slot = Arc::new(Mutex::new(None));
let record = Arc::new(EagerTensorRecord {
value,
key: key.clone(),
trace: None,
semantic_trace: semantic_trace.clone(),
requires_grad: false,
grad_slot: Arc::clone(&grad_slot),
ctx: Arc::clone(&ctx),
});
Self {
key,
trace: None,
semantic_trace,
requires_grad: false,
grad_slot,
ctx,
_record: record,
}
}
pub(crate) fn from_record(record: Arc<EagerTensorRecord>) -> Self {
Self {
key: record.key.clone(),
trace: record.trace.clone(),
semantic_trace: record.semantic_trace.clone(),
requires_grad: record.requires_grad,
grad_slot: Arc::clone(&record.grad_slot),
ctx: Arc::clone(&record.ctx),
_record: record,
}
}
pub fn detach(&self, engine: &mut dyn Engine) -> Result<Self> {
let tensor = self.duplicate_value(engine)?;
let semantic_trace = TracedTensor::from_tensor_symbolic_shape(tensor)?;
Ok(Self::new_untracked_value_record(
self.ctx.clone(),
Arc::clone(&self._record.value),
Some(semantic_trace),
))
}
pub fn detach_into(&self, engine: &mut dyn Engine, ctx: &Arc<EagerRuntime>) -> Result<Self> {
let tensor = self.to_tensor(&mut *engine)?;
Self::new_leaf_in_session(Arc::clone(ctx), tensor, false, engine)
}
pub fn value(&self) -> Result<ValueGuard<'_>> {
self._record.value.value("EagerTensor::value")
}
pub fn duplicate_value(&self, engine: &mut dyn Engine) -> Result<Tensor> {
self.duplicate_value_in_session(engine)
}
fn duplicate_host_value(&self) -> Option<Tensor> {
self.value().ok()?.duplicate_host_tensor().ok()
}
pub(crate) fn duplicate_value_in_session(&self, session: &mut dyn Engine) -> Result<Tensor> {
if let Some(tensor) = self.duplicate_host_value() {
return Ok(tensor);
}
let read = self
._record
.value
.tensor_read("EagerTensor::duplicate_value")?;
session.to_contiguous_read(read).map_err(Error::from)
}
#[allow(clippy::result_large_err)]
pub fn into_value(self) -> std::result::Result<Tensor, IntoValueError<Self>> {
if Arc::strong_count(&self._record) != 1 {
return Err(IntoValueError::NotUnique(self));
}
let Self { _record, .. } = self;
let record = match Arc::try_unwrap(_record) {
Ok(record) => record,
Err(record) => return Err(IntoValueError::NotUnique(Self::from_record(record))),
};
let EagerTensorRecord {
value,
key,
trace,
semantic_trace,
requires_grad,
grad_slot,
ctx,
} = record;
let value = match Arc::try_unwrap(value) {
Ok(value) => value,
Err(value) => {
let record = Arc::new(EagerTensorRecord {
value,
key,
trace,
semantic_trace,
requires_grad,
grad_slot,
ctx,
});
return Err(IntoValueError::NotUnique(Self::from_record(record)));
}
};
let AdValueRecord {
container,
dtype,
shape,
} = value;
let container = match Arc::try_unwrap(container) {
Ok(container) => container,
Err(container) => {
let record = Arc::new(EagerTensorRecord {
value: Arc::new(AdValueRecord {
container,
dtype,
shape,
}),
key,
trace,
semantic_trace,
requires_grad,
grad_slot,
ctx,
});
return Err(IntoValueError::NotUnique(Self::from_record(record)));
}
};
let (group, slot) = match container {
RetentionContainer::CallerOwned { tensor } => return Ok(*tensor),
RetentionContainer::Owned { tensor } => return Ok(tensor),
RetentionContainer::Pooled { group, slot } => (group, slot),
};
match group.into_tensor(slot) {
Ok(tensor) => Ok(tensor),
Err((group, error)) => {
let record = Arc::new(EagerTensorRecord {
value: Arc::new(AdValueRecord {
container: Arc::new(RetentionContainer::Pooled { group, slot }),
dtype,
shape,
}),
key,
trace,
semantic_trace,
requires_grad,
grad_slot,
ctx,
});
Err(IntoValueError::Extract {
value: Self::from_record(record),
error,
})
}
}
}
pub fn dtype(&self) -> DType {
self._record.value.dtype()
}
pub fn shape(&self) -> &[usize] {
self._record.value.shape()
}
pub fn tensor_read(&self) -> TensorRead<'_> {
self._record
.value
.tensor_read("EagerTensor::tensor_read")
.expect("validated eager value record")
}
pub fn to_tensor(&self, engine: &mut dyn Engine) -> Result<Tensor> {
self.duplicate_value(engine)
}
pub fn grad(&self) -> Result<Option<GradientValue>> {
self.grad_slot
.lock()
.map_err(|_| {
Error::runtime_state(
"eager_gradient_slot",
ErrorPhase::Execution,
"lock poisoned",
)
})
.map(|slot| {
slot.as_ref().map(|record| GradientValue {
record: Arc::clone(record),
})
})
}
pub fn clear_grad(&self) -> Result<()> {
*self.grad_slot.lock().map_err(|_| {
Error::runtime_state(
"eager_gradient_slot",
ErrorPhase::Execution,
"lock poisoned",
)
})? = None;
Ok(())
}
pub fn tracks_grad(&self) -> bool {
self.requires_grad
}
#[cfg(test)]
fn debug_trace_saved_value_count(&self) -> Option<usize> {
None
}
pub fn ctx_id(&self) -> ContextId {
self.ctx.id()
}
pub fn runtime(&self) -> &Arc<EagerRuntime> {
&self.ctx
}
pub fn same_context(&self, other: &Self) -> bool {
self.ctx_id() == other.ctx_id()
}
#[cfg(test)]
pub(crate) fn standard_graph_op(
engine: &mut dyn ExecutionEngine,
inputs: &[&Self],
build_graph: impl FnOnce(&[TensorInputKey]) -> Result<Arc<Graph<StdTensorOp>>>,
) -> Result<Vec<Self>> {
let Some(first) = inputs.first() else {
return Err(Error::Internal(
"standard eager graph op requires at least one input tensor".to_string(),
));
};
let ctx = Arc::clone(&first.ctx);
for tensor in inputs.iter().skip(1) {
if !first.same_context(tensor) {
return Err(Error::ContextMismatch {
lhs: first.ctx_id(),
rhs: tensor.ctx_id(),
});
}
}
let graph_input_keys = (0..inputs.len())
.map(|_| next_input_key())
.collect::<Vec<_>>();
let graph = build_graph(&graph_input_keys)?;
let initial_data = graph_input_keys
.iter()
.zip(inputs.iter())
.map(|(key, tensor)| {
Ok((
ValueKey::Input(key.clone()),
tensor.to_tensor(&mut *engine)?,
))
})
.collect::<Result<HashMap<_, _>>>()?;
let execution = ctx.exec_standard_graph_outputs(engine, graph.as_ref(), initial_data)?;
if execution.outputs.len() != graph.outputs().len() {
return Err(Error::Internal(format!(
"standard eager graph op expected {} graph outputs, got {}",
graph.outputs().len(),
execution.outputs.len()
)));
}
if !eager_grad_recording_enabled() || !inputs.iter().any(|input| input.requires_grad) {
return execution
.outputs
.into_iter()
.map(|output| {
Self::new_unregistered_result_with_semantic_trace(
Arc::clone(&ctx),
eager_val_key(),
output,
false,
None,
None,
)
})
.collect();
}
let recorded = record_eager_graph_outputs(
graph.as_ref(),
&graph_input_keys,
&execution.outputs,
inputs,
)?;
if recorded.traces.len() != execution.outputs.len() {
return Err(Error::Internal(format!(
"standard eager graph op expected {} eager traces, got {}",
execution.outputs.len(),
recorded.traces.len()
)));
}
recorded
.traces
.into_iter()
.zip(recorded.semantic_traces)
.zip(execution.outputs)
.map(|((trace, semantic_trace), output)| {
Self::new_result_with_semantic_trace(
Arc::clone(&ctx),
trace.key,
output,
trace.requires_grad,
trace.trace,
semantic_trace,
)
})
.collect()
}
pub fn backward(&self, engine: &mut dyn ExecutionEngine) -> Result<Gradients> {
if !self.shape().is_empty() {
return Err(Error::NonScalarGrad {
shape: self.shape().to_vec(),
});
}
let value = self.to_tensor(&mut *engine)?;
let seed = one_like_tensor(&value, &mut *engine)?;
self.backward_from_seed(engine, seed)
}
pub fn backward_with(
&self,
engine: &mut dyn ExecutionEngine,
cotangent: &EagerTensor,
) -> Result<Gradients> {
if !self.same_context(cotangent) {
return Err(Error::ContextMismatch {
lhs: self.ctx_id(),
rhs: cotangent.ctx_id(),
});
}
validate_seed_tensor("backward", self, cotangent)?;
let seed = cotangent.to_tensor(&mut *engine)?;
self.backward_from_seed(engine, seed)
}
fn backward_from_seed(
&self,
engine: &mut dyn ExecutionEngine,
seed: Tensor,
) -> Result<Gradients> {
let cotangent =
EagerTensor::new_result(Arc::clone(&self.ctx), eager_val_key(), seed, false, None)?;
let candidate_keys = {
let mut slots = self.ctx.lock_grad_slots()?;
let mut keys = Vec::new();
slots.retain(|key, slot| {
if slot.upgrade().is_some() {
keys.push(key.clone());
true
} else {
false
}
});
keys
};
let mut targets = Vec::new();
for key in candidate_keys {
let Some(record) = self.ctx.value_record(&key)? else {
continue;
};
if !record.requires_grad {
continue;
}
targets.push((key, EagerTensor::from_record(record)));
}
let wrts = targets.iter().map(|(_, wrt)| wrt).collect::<Vec<_>>();
let gradients = semantic_eager_vjp_many(&self.ctx, engine, self, &wrts, &cotangent)?;
let cotangents = (|| {
let session: &mut dyn Engine = &mut *engine;
let mut cotangents = HashMap::new();
for ((key, _), grad) in targets.into_iter().zip(gradients) {
let Some(grad) = grad else {
continue;
};
let tensor = match grad.into_value() {
Ok(tensor) => tensor,
Err(IntoValueError::NotUnique(handle)) => {
handle.duplicate_value_in_session(session)?
}
Err(IntoValueError::Extract { value, .. }) => {
value.duplicate_value_in_session(session)?
}
};
cotangents.insert(key, tensor);
}
self.ctx.store_grads(&cotangents, session)?;
Ok::<_, Error>(cotangents)
})()?;
Gradients::from_tensors(cotangents)
}
}
fn insert_pruning_dead<K: std::hash::Hash + Eq, V>(
map: &mut HashMap<K, Weak<V>>,
key: K,
value: Weak<V>,
) {
if map.len() == map.capacity() {
map.retain(|_, entry| entry.strong_count() > 0);
}
map.insert(key, value);
}
pub(crate) fn eager_val_key() -> ValueKey<StdTensorOp> {
ValueKey::Input(next_input_key())
}
pub(crate) struct RecordedEagerTrace {
pub(crate) key: ValueKey<StdTensorOp>,
pub(crate) trace: Option<EagerTrace>,
pub(crate) requires_grad: bool,
}
pub(crate) struct RecordedEagerOutputs {
pub(crate) traces: Vec<RecordedEagerTrace>,
pub(crate) semantic_traces: Vec<Option<TracedTensor>>,
}
pub(crate) fn record_eager_outputs_in_session(
op: &StdTensorOp,
outputs: &[&Tensor],
inputs: &[&EagerTensor],
session: &mut dyn Engine,
) -> Result<RecordedEagerOutputs> {
let metadata = outputs
.iter()
.map(|output| tensor_meta_from_tensor(output))
.collect();
record_eager_outputs_inner(op, metadata, inputs, Some(session))
}
pub(crate) fn record_eager_value_outputs_in_session(
op: &StdTensorOp,
outputs: &[&TensorValue],
inputs: &[&EagerTensor],
session: &mut dyn Engine,
) -> Result<RecordedEagerOutputs> {
let metadata = outputs
.iter()
.map(|output| tensor_meta_from_value(output))
.collect();
record_eager_outputs_inner(op, metadata, inputs, Some(session))
}
fn record_eager_outputs_inner(
op: &StdTensorOp,
output_metadata: Vec<TensorMeta>,
inputs: &[&EagerTensor],
session: Option<&mut dyn Engine>,
) -> Result<RecordedEagerOutputs> {
let semantic_traces = record_semantic_eager_outputs(op, &output_metadata, inputs, session)?;
record_eager_outputs_from_metadata(output_metadata, semantic_traces, inputs)
}
fn record_semantic_eager_outputs(
op: &StdTensorOp,
output_metadata: &[TensorMeta],
inputs: &[&EagerTensor],
mut session: Option<&mut dyn Engine>,
) -> Result<Vec<Option<TracedTensor>>> {
let mut owned_constants = Vec::<TracedTensor>::new();
for input in inputs {
if input.semantic_trace.is_none() {
let value = match session.as_deref_mut() {
Some(session) => input.duplicate_value_in_session(session)?,
None => input.duplicate_host_value().ok_or_else(|| {
Error::unsupported(
"record_semantic_eager_outputs",
ErrorPhase::Execution,
"a non-host untracked input needs an engine to become a constant leaf",
)
})?,
};
owned_constants.push(TracedTensor::from_tensor_symbolic_shape(value)?);
}
}
let mut constants = owned_constants.iter();
let mut semantic_inputs: Vec<&TracedTensor> = inputs
.iter()
.map(|input| {
input
.semantic_trace
.as_ref()
.unwrap_or_else(|| constants.next().expect("materialized constant"))
})
.collect();
let promotion_plan =
eager_input_promotion_plan(op, inputs.len(), |index| inputs[index].dtype());
let promoted_semantic_inputs = if semantic_inputs.iter().enumerate().any(|(index, semantic)| {
semantic.dtype != promotion_plan.target_dtype(index, semantic.dtype)
}) {
Some(
semantic_inputs
.iter()
.enumerate()
.map(|(index, &semantic)| {
let target = promotion_plan.target_dtype(index, semantic.dtype);
if semantic.dtype == target {
Ok(Cow::Borrowed(semantic))
} else {
semantic.cast(target).map(Cow::Owned)
}
})
.collect::<Result<Vec<Cow<'_, TracedTensor>>>>()?,
)
} else {
None
};
if let Some(promoted_semantic_inputs) = &promoted_semantic_inputs {
semantic_inputs = promoted_semantic_inputs.iter().map(Cow::as_ref).collect();
}
let exact_semantic_inputs = if matches!(op, StdTensorOp::Concatenate { .. }) {
Some(
semantic_inputs
.iter()
.zip(inputs)
.map(|(&semantic, input)| {
if semantic.is_concrete_shape() {
Ok(semantic.clone())
} else {
semantic.reshape(input.shape())
}
})
.collect::<Result<Vec<_>>>()?,
)
} else {
None
};
if let Some(exact_semantic_inputs) = &exact_semantic_inputs {
semantic_inputs = exact_semantic_inputs.iter().collect();
}
let outputs = tenferro_runtime::extension::append_raw_eager_outputs(
op.clone(),
&semantic_inputs,
output_metadata,
)?;
Ok(outputs.into_iter().map(Some).collect())
}
#[cfg(test)]
fn record_eager_graph_outputs(
graph: &Graph<StdTensorOp>,
graph_input_keys: &[TensorInputKey],
outputs: &[Tensor],
inputs: &[&EagerTensor],
) -> Result<RecordedEagerOutputs> {
let semantic_traces = record_semantic_eager_graph_outputs(graph, graph_input_keys, inputs)?;
let output_metadata = outputs.iter().map(tensor_meta_from_tensor);
record_eager_outputs_from_metadata(output_metadata, semantic_traces, inputs)
}
#[cfg(test)]
fn record_semantic_eager_graph_outputs(
graph: &Graph<StdTensorOp>,
graph_input_keys: &[TensorInputKey],
inputs: &[&EagerTensor],
) -> Result<Vec<Option<TracedTensor>>> {
let Some(semantic_inputs) = inputs
.iter()
.map(|input| input.semantic_trace.as_ref())
.collect::<Option<Vec<_>>>()
else {
return Ok(vec![None; graph.outputs().len()]);
};
if graph_input_keys.len() != semantic_inputs.len() {
return Err(Error::Internal(format!(
"semantic graph recording expected {} input keys, got {}",
semantic_inputs.len(),
graph_input_keys.len()
)));
}
let mut values = HashMap::new();
for (key, tensor) in graph_input_keys.iter().zip(semantic_inputs) {
values.insert(ValueKey::Input(key.clone()), tensor.clone());
}
for op_node in graph.operations() {
let input_values = op_node
.inputs
.iter()
.map(|input| {
let key = match input {
ValueRef::Local(local_id) => &graph.values()[*local_id].key,
ValueRef::External(key) => key,
};
values.get(key).cloned().ok_or_else(|| {
Error::Internal(format!(
"semantic graph recording missing value for {key:?}"
))
})
})
.collect::<Result<Vec<_>>>()?;
let input_refs = input_values.iter().collect::<Vec<_>>();
let semantic_outputs = match &op_node.operation {
StdTensorOp::Extension(ext) => {
tenferro_runtime::extension::apply(Arc::clone(ext), &input_refs)?
}
op => tenferro_runtime::extension::apply_standard_op(op.clone(), &input_refs)?,
};
if semantic_outputs.len() != op_node.outputs.len() {
return Err(Error::Internal(format!(
"semantic graph recording expected {} outputs for {:?}, got {}",
op_node.outputs.len(),
op_node.operation,
semantic_outputs.len()
)));
}
for (output_id, output) in op_node.outputs.iter().copied().zip(semantic_outputs) {
values.insert(graph.values()[output_id].key.clone(), output);
}
}
graph
.outputs()
.iter()
.map(|&output_id| {
let key = &graph.values()[output_id].key;
values.get(key).cloned().map(Some).ok_or_else(|| {
Error::Internal(format!(
"semantic graph recording missing output for {key:?}"
))
})
})
.collect()
}
fn record_eager_outputs_from_metadata(
output_metadata: impl IntoIterator<Item = TensorMeta>,
semantic_traces: Vec<Option<TracedTensor>>,
inputs: &[&EagerTensor],
) -> Result<RecordedEagerOutputs> {
let output_metadata = output_metadata.into_iter().collect::<Vec<_>>();
if semantic_traces.len() != output_metadata.len() {
return Err(Error::Internal(format!(
"eager recording expected {} semantic traces, got {}",
output_metadata.len(),
semantic_traces.len()
)));
}
let requires_grad =
eager_grad_recording_enabled() && inputs.iter().any(|input| input.requires_grad);
let trace_count = output_metadata.len();
let residual_trace = EagerTrace::new(inputs);
let traces = (0..trace_count)
.map(|_| RecordedEagerTrace {
key: eager_val_key(),
trace: Some(residual_trace.clone()),
requires_grad,
})
.collect();
Ok(RecordedEagerOutputs {
traces,
semantic_traces,
})
}
fn tensor_meta_from_value(value: &TensorValue) -> TensorMeta {
TensorMeta::exact(
value.dtype(),
value.shape().iter().copied().map(SymDim::from).collect(),
)
}
#[cfg(test)]
pub(crate) fn zero_like_tensor(input: &Tensor, backend: &mut dyn Engine) -> Result<Tensor> {
let host = match input.dtype() {
DType::External(type_id) => {
return Err(Error::unsupported(
"zero_like_tensor",
ErrorPhase::GraphBuild,
format!(
"an externally defined payload ({:?}) has no zero-like runtime tensor",
DType::External(type_id)
),
));
}
DType::F32 => Tensor::from_typed::<f32>(TypedTensor::zeros(input.shape().to_vec())?),
DType::F64 => Tensor::from_typed::<f64>(TypedTensor::zeros(input.shape().to_vec())?),
DType::I32 => Tensor::from_typed::<i32>(TypedTensor::zeros(input.shape().to_vec())?),
DType::I64 => Tensor::from_typed::<i64>(TypedTensor::zeros(input.shape().to_vec())?),
DType::Bool => Tensor::from_typed::<bool>(TypedTensor::from_vec_col_major(
input.shape().to_vec(),
vec![false; input.shape().iter().product()],
)?),
DType::C32 => Tensor::from_typed::<tenferro_tensor::Complex32>(TypedTensor::zeros(
input.shape().to_vec(),
)?),
DType::C64 => Tensor::from_typed::<tenferro_tensor::Complex64>(TypedTensor::zeros(
input.shape().to_vec(),
)?),
};
backend
.upload_host_tensor(TensorRead::from_tensor(&host))
.map_err(Error::from)
}
pub(crate) fn one_like_tensor(input: &Tensor, session: &mut dyn Engine) -> Result<Tensor> {
let host = ones_tensor(input.dtype(), input.shape().to_vec())?;
session
.upload_host_tensor(TensorRead::from_tensor(&host))
.map_err(Error::from)
}
#[cfg(test)]
mod tests;