tenferro_ad/eager.rs
1use std::borrow::Cow;
2use std::cell::{Cell, RefCell};
3use std::cmp::Reverse;
4use std::collections::HashMap;
5use std::env;
6use std::fmt;
7use std::marker::PhantomData;
8use std::mem::{size_of, size_of_val};
9use std::rc::Rc;
10#[cfg(test)]
11use std::sync::atomic::{AtomicUsize, Ordering};
12use std::sync::{Arc, Mutex, MutexGuard, OnceLock, Weak};
13use std::time::{Duration, Instant};
14
15use lru::LruCache;
16use num_complex::Complex64;
17
18use crate::extension::{
19 validate_eager_extension_target, EagerExtensionBackendKind, EagerExtensionTarget,
20};
21
22#[cfg(test)]
23use computegraph::graph::Graph;
24use computegraph::ValueKey;
25#[cfg(test)]
26use computegraph::ValueRef;
27#[cfg(test)]
28use tenferro_ops::input_key::TensorInputKey;
29use tenferro_ops::{std_tensor_op::StdTensorOp, SymDim, TensorMeta};
30use tenferro_runtime::ad_support::{
31 analyze_deferred_semantic_trace, compile_ad_source, ones_tensor, RetainedValue,
32};
33use tenferro_runtime::program::{ProgramValueMetadata, SemanticFingerprint, SemanticProgram};
34use tenferro_runtime::{
35 CompiledGraph, EngineId, EngineTarget, ErrorPhase, ExecutionContextIdentity, ExecutionEngine,
36 ExtensionModule, GraphCompiler, Prepared, Runtime, RuntimeConfigError, RuntimeEpoch,
37 TracedTensor,
38};
39use tenferro_tensor::Engine;
40#[cfg(test)]
41use tenferro_tensor::TypedTensor;
42use tenferro_tensor::{
43 AllocationGroup, CacheStats, CompareDir, DType, DescriptorSlot, DotGeneralConfig, GatherConfig,
44 GroupError, IntoShapeVec, PadConfig, ScatterConfig, SliceConfig, Tensor, TensorRead,
45 TensorScalar, TensorValue, TensorView,
46};
47
48use crate::eager_backend::{
49 cpu_host_owned_read, engine_kind_registration, unregisterable_engine_error,
50};
51#[cfg(test)]
52use crate::eager_exec::exec_standard_op_on_tensor_reads_in_session;
53use crate::eager_exec::{eager_input_promotion_plan, exec_extension_op_on_tensor_reads};
54use crate::error::{ContextId, Error, Result};
55use crate::metadata::tensor_meta_from_tensor;
56use crate::semantic_extension::SemanticExtensionRuleSet;
57use crate::traced::{derivative_trace_from_frozen_program, next_input_key};
58use crate::transform_cache::{AdTransformCache, AdTransformCacheLimits};
59
60use crate::AdContext;
61
62pub(crate) type GradSlot = Arc<Mutex<Option<Arc<AdValueRecord>>>>;
63pub(crate) type WeakGradSlot = Weak<Mutex<Option<Arc<AdValueRecord>>>>;
64
65mod composite;
66mod residuals;
67pub(crate) use residuals::{finish_residuals, EagerTrace};
68
69#[derive(Debug, Default, Clone)]
70struct EagerOpProfileEntry {
71 calls: usize,
72 total_time: Duration,
73}
74
75thread_local! {
76 static EAGER_OP_PROFILE_STATE: RefCell<HashMap<&'static str, EagerOpProfileEntry>> =
77 RefCell::new(HashMap::new());
78 static EAGER_NO_GRAD_DEPTH: Cell<usize> = const { Cell::new(0) };
79 static EAGER_CAPTURE_DEPTH: Cell<usize> = const { Cell::new(0) };
80 #[cfg(test)]
81 static EAGER_OP_PROFILE_ENABLED_OVERRIDE: RefCell<Option<bool>> = const { RefCell::new(None) };
82 #[cfg(test)]
83 static EAGER_OP_PROFILE_PRINT_EVERY_OVERRIDE: RefCell<Option<Option<usize>>> = const { RefCell::new(None) };
84}
85
86#[cfg(test)]
87pub(crate) static EAGER_SEMANTIC_VJP_EXECUTIONS: AtomicUsize = AtomicUsize::new(0);
88
89pub(crate) fn eager_grad_recording_enabled() -> bool {
90 EAGER_NO_GRAD_DEPTH.with(|depth| depth.get() == 0)
91}
92
93pub(crate) fn eager_capture_active() -> bool {
94 EAGER_CAPTURE_DEPTH.with(|depth| depth.get() > 0)
95}
96
97/// Scope guard that temporarily disables eager operation recording.
98///
99/// Values computed while this guard is alive are concrete eager tensors, but
100/// they do not participate in reverse-mode gradient tracking.
101///
102/// # Examples
103///
104/// ```
105/// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
106///
107/// let ctx = EagerRuntime::new()?;
108/// let x = EagerTensor::requires_grad_in(
109/// Tensor::from_vec_col_major(vec![1], vec![2.0_f64]).unwrap(),
110/// ctx.clone(),
111/// )?;
112/// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
113/// let _guard = ctx.no_grad();
114/// s.mul(&x, &x)
115/// })?;
116/// assert!(!y.tracks_grad());
117/// # Ok::<(), tenferro_ad::Error>(())
118/// ```
119#[derive(Debug)]
120pub struct EagerNoGradGuard {
121 active: bool,
122 // Thread-local depth guard: must not be Send so it cannot be moved to and
123 // dropped on another thread (which would corrupt the creator's depth).
124 _not_send: PhantomData<Rc<()>>,
125}
126
127impl Drop for EagerNoGradGuard {
128 fn drop(&mut self) {
129 if !self.active {
130 return;
131 }
132 EAGER_NO_GRAD_DEPTH.with(|depth| {
133 depth.set(depth.get().saturating_sub(1));
134 });
135 self.active = false;
136 }
137}
138
139/// Scope guard that keeps semantic-trace recording active for untracked
140/// intermediates.
141///
142/// Under active-edge semantics (issue #1665 Def 1), an operation whose inputs
143/// are all untracked produces no autograd nodes and drops its semantic trace.
144/// Inside this guard, such operations still record their semantic trace, so a
145/// later functional JVP/VJP can differentiate with respect to an untracked or
146/// detached leaf. This replaces the pre-Def-1 implicit recording.
147///
148/// # Examples
149///
150/// ```
151/// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
152///
153/// let ctx = EagerRuntime::new()?;
154/// let x = EagerTensor::from_tensor_in(
155/// Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0]).unwrap(),
156/// ctx.clone(),
157/// )?;
158/// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
159/// let _capture = ctx.capture_trace();
160/// s.mul(&x, &x)
161/// })?;
162/// let seed = EagerTensor::from_tensor_in(
163/// Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 1.0]).unwrap(),
164/// ctx.clone(),
165/// )?;
166/// let dx = ctx.vjp(&mut tenferro_cpu::CpuEngine::serial(), &y, &x, &seed)?;
167/// assert_eq!(dx.value()?.as_slice::<f64>().unwrap(), &[2.0, 4.0]);
168/// # Ok::<(), tenferro_ad::Error>(())
169/// ```
170#[derive(Debug)]
171pub struct EagerTraceCaptureGuard {
172 active: bool,
173 // Thread-local depth guard: must not be Send so it cannot be moved to and
174 // dropped on another thread (which would corrupt the creator's depth).
175 _not_send: PhantomData<Rc<()>>,
176}
177
178impl Drop for EagerTraceCaptureGuard {
179 fn drop(&mut self) {
180 if !self.active {
181 return;
182 }
183 EAGER_CAPTURE_DEPTH.with(|depth| {
184 depth.set(depth.get().saturating_sub(1));
185 });
186 self.active = false;
187 }
188}
189
190pub(crate) fn eager_op_profile_enabled() -> bool {
191 #[cfg(test)]
192 if let Some(value) = EAGER_OP_PROFILE_ENABLED_OVERRIDE.with(|state| *state.borrow()) {
193 return value;
194 }
195
196 static ENABLED: OnceLock<bool> = OnceLock::new();
197 *ENABLED.get_or_init(|| env::var("TENFERRO_PROFILE_EAGER_OP_AGG").is_ok())
198}
199
200pub(crate) fn eager_op_profile_start() -> Option<Instant> {
201 eager_op_profile_enabled().then(Instant::now)
202}
203
204pub(crate) fn record_eager_op_profile(section: &'static str, elapsed: Duration) {
205 if !eager_op_profile_enabled() {
206 return;
207 }
208 EAGER_OP_PROFILE_STATE.with(|state| {
209 let mut state = state.borrow_mut();
210 let entry = state.entry(section).or_default();
211 entry.calls += 1;
212 entry.total_time += elapsed;
213 });
214}
215
216pub(crate) fn profile_eager_op_section<T>(section: &'static str, f: impl FnOnce() -> T) -> T {
217 if !eager_op_profile_enabled() {
218 return f();
219 }
220 let started = Instant::now();
221 let result = f();
222 record_eager_op_profile(section, started.elapsed());
223 result
224}
225
226pub(crate) fn maybe_print_eager_op_profile() {
227 if !eager_op_profile_enabled() {
228 return;
229 }
230 let Some(print_every) = eager_op_profile_print_every() else {
231 return;
232 };
233 if print_every == 0 {
234 return;
235 }
236
237 let should_print = EAGER_OP_PROFILE_STATE.with(|state| {
238 state
239 .borrow()
240 .get("nary_op.total")
241 .is_some_and(|entry| entry.calls % print_every == 0)
242 });
243 if should_print {
244 print_and_reset_eager_op_profile();
245 }
246}
247
248fn eager_op_profile_print_every() -> Option<usize> {
249 #[cfg(test)]
250 if let Some(value) = EAGER_OP_PROFILE_PRINT_EVERY_OVERRIDE.with(|state| *state.borrow()) {
251 return value;
252 }
253
254 env::var("TENFERRO_PROFILE_EAGER_OP_PRINT_EVERY")
255 .ok()?
256 .parse()
257 .ok()
258}
259
260pub(crate) fn print_and_reset_eager_op_profile() {
261 EAGER_OP_PROFILE_STATE.with(|state| {
262 let mut entries: Vec<_> = state
263 .borrow()
264 .iter()
265 .map(|(section, entry)| (*section, entry.clone()))
266 .collect();
267 state.borrow_mut().clear();
268 entries.sort_by_key(|(_, entry)| Reverse(entry.total_time));
269
270 eprintln!("=== tenferro eager op profile ===");
271 for (section, entry) in entries {
272 let Some(per_call_us) = eager_op_profile_per_call_us(&entry) else {
273 continue;
274 };
275 eprintln!(
276 "{section}: calls={} total={:.6}ms per_call={:.3}us",
277 entry.calls,
278 entry.total_time.as_secs_f64() * 1.0e3,
279 per_call_us,
280 );
281 }
282 });
283}
284
285fn eager_op_profile_per_call_us(entry: &EagerOpProfileEntry) -> Option<f64> {
286 (entry.calls != 0).then(|| entry.total_time.as_secs_f64() * 1.0e6 / entry.calls as f64)
287}
288
289fn runtime_config_error(op: &'static str, source: RuntimeConfigError) -> Error {
290 Error::runtime_state_source(op, ErrorPhase::Execution, source)
291}
292
293fn runtime_state_source<E>(op: &'static str, source: E) -> Error
294where
295 E: std::error::Error + Send + Sync + 'static,
296{
297 Error::runtime_state_source(op, ErrorPhase::Execution, source)
298}
299
300/// Run `prepared` on `engine` and return owned compact outputs.
301pub(crate) fn run_prepared_on(
302 runtime: &Runtime,
303 engine: &mut dyn ExecutionEngine,
304 prepared: &Prepared,
305 inputs: &[TensorRead<'_>],
306) -> Result<Vec<Tensor>> {
307 let outputs = runtime.run(&mut [&mut *engine], prepared, inputs)?;
308 outputs
309 .into_iter()
310 .map(|output| output.into_tensor(&mut *engine).map_err(Error::from))
311 .collect()
312}
313
314/// Run a session-capable prepared extension executor on `engine`, lending it
315/// the engine's extension caches for the call.
316pub(crate) fn execute_prepared_extension_on(
317 engine: &mut dyn ExecutionEngine,
318 executor: &dyn tenferro_runtime::PreparedOperationExecutor,
319 inputs: &[TensorRead<'_>],
320) -> Result<Vec<Tensor>> {
321 with_lent_extension_caches(engine, |engine, caches| {
322 executor.execute_in_session(engine, caches, inputs)
323 })
324}
325
326/// Lend `engine`'s extension caches to `callback` and restore them on every
327/// exit path, including unwind.
328///
329/// A callback panic also marks the engine broken before it is resumed: the
330/// caches may be half updated, so the engine must not accept more work.
331pub(crate) fn with_lent_extension_caches<R>(
332 engine: &mut dyn ExecutionEngine,
333 callback: impl FnOnce(&mut dyn ExecutionEngine, &mut tenferro_runtime::ExtensionCacheStore) -> R,
334) -> R {
335 let mut caches = engine.take_extension_caches();
336 let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
337 callback(&mut *engine, &mut caches)
338 }));
339 engine.restore_extension_caches(caches);
340 match result {
341 Ok(value) => value,
342 Err(payload) => {
343 engine.mark_broken();
344 std::panic::resume_unwind(payload)
345 }
346 }
347}
348
349/// The eager extension backend kind of an engine target, or `None` for an
350/// engine type without eager extension support.
351pub(crate) fn eager_extension_backend_kind(
352 target: &EngineTarget,
353) -> Option<EagerExtensionBackendKind> {
354 if target.context_identity() == ExecutionContextIdentity::of::<tenferro_cpu::CpuEngine>() {
355 return Some(EagerExtensionBackendKind::Cpu);
356 }
357 #[cfg(feature = "cuda")]
358 if target.context_identity() == ExecutionContextIdentity::of::<tenferro_gpu::cuda::CudaEngine>()
359 {
360 return Some(EagerExtensionBackendKind::Cuda);
361 }
362 None
363}
364
365/// Stats for caches owned by an [`EagerRuntime`].
366///
367/// `retained_bytes` fields are logical payload estimates, not process RSS.
368#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
369pub struct EagerRuntimeCacheStats {
370 /// Eager AD transform memoization cache.
371 pub ad_transforms: CacheStats,
372 /// Prepared eager derivative program cache.
373 pub prepared_derivatives: CacheStats,
374}
375
376#[cfg(test)]
377pub(crate) struct EagerGraphExecution {
378 pub(crate) outputs: Vec<Tensor>,
379}
380
381/// A read-only value view retained by an eager tensor record.
382///
383/// The guard borrows the record's allocation group. It never owns a tensor and
384/// cannot be converted into a mutable view.
385///
386/// # Examples
387///
388/// ```
389/// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
390///
391/// let ctx = EagerRuntime::new()?;
392/// let value = EagerTensor::from_tensor_in(
393/// Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0])?,
394/// ctx,
395/// )?;
396/// let view = value.value()?;
397/// assert_eq!(view.shape(), &[2]);
398/// # Ok::<(), tenferro_ad::Error>(())
399/// ```
400#[derive(Debug)]
401pub struct ValueGuard<'a> {
402 view: TensorView<'a>,
403}
404
405impl<'a> ValueGuard<'a> {
406 /// Return the scalar dtype of the retained value.
407 pub fn dtype(&self) -> DType {
408 self.view.dtype()
409 }
410
411 /// Return the logical shape of the retained value.
412 pub fn shape(&self) -> &[usize] {
413 self.view.shape()
414 }
415
416 /// Borrow the dtype-erased tensor view.
417 pub fn as_tensor_view(&self) -> &TensorView<'_> {
418 &self.view
419 }
420
421 /// Borrow compact host bytes through the tensor's explicit scalar type.
422 ///
423 /// Backend-resident values return the backend's typed host-access error;
424 /// this method does not download storage implicitly.
425 ///
426 /// # Errors
427 ///
428 /// Returns [`tenferro_tensor::ValidationError::DTypeMismatch`] when
429 /// `T` does not match the view dtype, [`tenferro_tensor::ValidationError::NonContiguousViewAsSlice`]
430 /// for a non-contiguous view, or [`tenferro_tensor::Error::HostAccess`]
431 /// when backend storage cannot be mapped as a host slice.
432 pub fn as_slice<T: TensorScalar>(&self) -> tenferro_tensor::Result<&'a [T]> {
433 self.view.as_slice()
434 }
435
436 fn duplicate_host_tensor(&self) -> tenferro_tensor::Result<Tensor> {
437 match &self.view {
438 TensorView::F32(view) => {
439 <f32 as TensorScalar>::into_tensor(view.shape().to_vec(), view.as_slice()?.to_vec())
440 }
441 TensorView::F64(view) => {
442 <f64 as TensorScalar>::into_tensor(view.shape().to_vec(), view.as_slice()?.to_vec())
443 }
444 TensorView::I32(view) => {
445 <i32 as TensorScalar>::into_tensor(view.shape().to_vec(), view.as_slice()?.to_vec())
446 }
447 TensorView::I64(view) => {
448 <i64 as TensorScalar>::into_tensor(view.shape().to_vec(), view.as_slice()?.to_vec())
449 }
450 TensorView::Bool(view) => <bool as TensorScalar>::into_tensor(
451 view.shape().to_vec(),
452 view.as_slice()?.to_vec(),
453 ),
454 TensorView::C32(view) => <num_complex::Complex32 as TensorScalar>::into_tensor(
455 view.shape().to_vec(),
456 view.as_slice()?.to_vec(),
457 ),
458 TensorView::C64(view) => <num_complex::Complex64 as TensorScalar>::into_tensor(
459 view.shape().to_vec(),
460 view.as_slice()?.to_vec(),
461 ),
462 }
463 }
464}
465
466/// Read-only retained gradient value.
467///
468/// # Examples
469///
470/// ```
471/// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
472///
473/// let ctx = EagerRuntime::new()?;
474/// let x = EagerTensor::requires_grad_in(
475/// Tensor::from_vec_col_major(vec![2], vec![2.0_f64, 3.0])?,
476/// ctx,
477/// )?;
478/// let loss = x.runtime().with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
479/// let squared = s.mul(&x, &x)?;
480/// s.reduce_sum(&squared, Some(&[0]))
481/// })?;
482/// let _gradients = loss.backward(&mut tenferro_cpu::CpuEngine::serial())?;
483/// let gradient = x.grad()?.expect("tracked leaf has a gradient");
484/// assert_eq!(gradient.shape(), &[2]);
485/// # Ok::<(), tenferro_ad::Error>(())
486/// ```
487#[derive(Clone, Debug)]
488pub struct GradientValue {
489 record: Arc<AdValueRecord>,
490}
491
492impl GradientValue {
493 /// Return the scalar dtype of the gradient.
494 pub fn dtype(&self) -> DType {
495 self.record.dtype()
496 }
497
498 /// Return the logical shape of the gradient.
499 pub fn shape(&self) -> &[usize] {
500 self.record.shape()
501 }
502
503 /// Borrow the gradient's value guard.
504 ///
505 /// # Errors
506 ///
507 /// Returns [`Error::RuntimeState`] when the retained gradient record is
508 /// unavailable or its allocation-group descriptor is invalid.
509 pub fn value(&self) -> Result<ValueGuard<'_>> {
510 self.record.value("GradientValue::value")
511 }
512
513 /// Borrow the gradient as a dtype-erased read target.
514 ///
515 /// # Errors
516 ///
517 /// Returns [`Error::RuntimeState`] when the retained gradient record or
518 /// its allocation-group descriptor is unavailable.
519 pub fn tensor_read(&self) -> Result<TensorRead<'_>> {
520 self.record.tensor_read("GradientValue::tensor_read")
521 }
522
523 /// Borrow a compact host slice without downloading backend storage.
524 ///
525 /// # Errors
526 ///
527 /// Returns [`Error::RuntimeState`] when the retained value is unavailable,
528 /// [`tenferro_tensor::ValidationError::DTypeMismatch`] when `T` does
529 /// not match the gradient dtype, or [`tenferro_tensor::Error::HostAccess`]
530 /// when backend storage cannot be mapped as a host slice.
531 pub fn as_slice<T: TensorScalar>(&self) -> tenferro_tensor::Result<&[T]> {
532 self.record
533 .value("GradientValue::as_slice")
534 .map_err(|error| {
535 tenferro_tensor::Error::runtime_state_source("GradientValue::as_slice", error)
536 })?
537 .as_slice()
538 }
539
540 /// Explicitly copy a host-resident gradient into a standalone tensor.
541 ///
542 /// # Errors
543 ///
544 /// Returns [`Error::RuntimeState`] when the retained value or execution
545 /// session is unavailable, or a typed backend/host-access error when the
546 /// gradient cannot be materialized as a contiguous tensor.
547 pub fn to_tensor(&self, engine: &mut dyn Engine) -> Result<Tensor> {
548 let value = self
549 .record
550 .value("GradientValue::to_tensor")
551 .map_err(|error| {
552 Error::runtime_state_source(
553 "GradientValue::to_tensor",
554 ErrorPhase::Execution,
555 error,
556 )
557 })?;
558 match value.duplicate_host_tensor() {
559 Ok(tensor) => Ok(tensor),
560 Err(_) => {
561 let read = self.record.tensor_read("GradientValue::to_tensor")?;
562 engine.to_contiguous_read(read).map_err(Error::from)
563 }
564 }
565 }
566}
567
568/// Move-only accumulated gradient bundle backed by one allocation group.
569///
570/// # Examples
571///
572/// ```
573/// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
574///
575/// let ctx = EagerRuntime::new()?;
576/// let x = EagerTensor::requires_grad_in(
577/// Tensor::from_vec_col_major(vec![2], vec![2.0_f64, 3.0])?,
578/// ctx,
579/// )?;
580/// let loss = x.runtime().with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
581/// let squared = s.mul(&x, &x)?;
582/// s.reduce_sum(&squared, Some(&[0]))
583/// })?;
584/// let gradients = loss.backward(&mut tenferro_cpu::CpuEngine::serial())?;
585/// assert!(!gradients.is_empty());
586/// # Ok::<(), tenferro_ad::Error>(())
587/// ```
588#[derive(Debug)]
589pub struct Gradients {
590 group: AllocationGroup,
591 slots: HashMap<ValueKey<StdTensorOp>, DescriptorSlot>,
592}
593
594impl Gradients {
595 fn from_tensors(tensors: HashMap<ValueKey<StdTensorOp>, Tensor>) -> Result<Self> {
596 if tensors.is_empty() {
597 return Err(Error::runtime_state(
598 "backward",
599 ErrorPhase::Execution,
600 "no gradients were produced for reachable tracked leaves; use EagerRuntime::grad_optional or vjp_optional to observe inactive differentiation",
601 ));
602 }
603 let (keys, values): (Vec<_>, Vec<_>) = tensors.into_iter().unzip();
604 let (group, bindings) = AllocationGroup::from_tensors(values).map_err(|error| {
605 Error::runtime_state_source("Gradients::from_tensors", ErrorPhase::Execution, error)
606 })?;
607 let slots = keys.into_iter().zip(bindings).collect();
608 Ok(Self { group, slots })
609 }
610
611 /// Return the number of retained gradient descriptors.
612 pub fn len(&self) -> usize {
613 self.slots.len()
614 }
615
616 /// Return whether no gradient was produced.
617 pub fn is_empty(&self) -> bool {
618 self.slots.is_empty()
619 }
620
621 /// Borrow one gradient view by its local value key.
622 pub fn grad(&self, key: &ValueKey<StdTensorOp>) -> Option<TensorView<'_>> {
623 let slot = self.slots.get(key).copied()?;
624 let mut reads = self.group.read_views(std::slice::from_ref(&slot)).ok()?;
625 match reads.pop()? {
626 TensorRead::View(view) => Some(view),
627 TensorRead::Tensor(_) => None,
628 }
629 }
630
631 /// Consume one gradient owner while leaving the bundle unchanged on failure.
632 ///
633 /// # Errors
634 ///
635 /// Returns [`tenferro_tensor::Error::RuntimeState`] when the descriptor is
636 /// invalid or its allocation is aliased. A missing key is reported as
637 /// `Ok(None)`.
638 pub fn take_grad(
639 &mut self,
640 key: &ValueKey<StdTensorOp>,
641 ) -> tenferro_tensor::Result<Option<Tensor>> {
642 let Some(&slot) = self.slots.get(key) else {
643 return Ok(None);
644 };
645 let tensor = self.group.take_tensor(slot).map_err(|error| {
646 tenferro_tensor::Error::runtime_state_source("Gradients::take_grad", error)
647 })?;
648 self.slots.remove(key);
649 Ok(Some(tensor))
650 }
651}
652
653/// Error returned when a value cannot be consumed without changing its owner.
654///
655/// # Examples
656///
657/// ```
658/// use tenferro_ad::{EagerRuntime, EagerTensor, IntoValueError, Tensor};
659///
660/// let ctx = EagerRuntime::new()?;
661/// let value = EagerTensor::from_tensor_in(
662/// Tensor::from_vec_col_major(vec![1], vec![1.0_f64])?,
663/// ctx,
664/// )?;
665/// let _shared = value.clone();
666/// assert!(matches!(
667/// value.into_value(),
668/// Err(IntoValueError::NotUnique(_))
669/// ));
670/// # Ok::<(), tenferro_ad::Error>(())
671/// ```
672#[derive(Debug)]
673pub enum IntoValueError<H> {
674 /// Another eager handle, tape record, or checkpoint retains the value.
675 NotUnique(H),
676 /// Group extraction failed after the handle was uniquely acquired.
677 Extract { value: H, error: GroupError },
678}
679
680impl<H> std::fmt::Display for IntoValueError<H> {
681 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
682 match self {
683 Self::NotUnique(_) => formatter.write_str("eager value is retained by another handle"),
684 Self::Extract { error, .. } => {
685 write!(formatter, "eager value extraction failed: {error}")
686 }
687 }
688 }
689}
690
691impl<H: std::fmt::Debug + Send + Sync + 'static> std::error::Error for IntoValueError<H> {}
692
693/// What one direct retention container holds.
694// The pooled variant owns an allocation group inline; boxing it would add an
695// allocation to every retained value on the hot path.
696#[allow(clippy::large_enum_variant)]
697#[derive(Debug)]
698enum RetentionContainer {
699 /// Pooled storage owned through an allocation group and one descriptor slot.
700 Pooled {
701 group: AllocationGroup,
702 slot: DescriptorSlot,
703 },
704 /// A caller-owned value the runtime retains without taking pool ownership.
705 ///
706 /// The value returns to its owner when the record drops, and the runtime never
707 /// substitutes it for pooled storage. A caller-owned payload has no typed
708 /// descriptor view, so only its read and consume paths are available.
709 CallerOwned {
710 /// Boxed because a tensor value is much larger than the pooled variant.
711 tensor: Box<Tensor>,
712 },
713 /// An untracked result held directly.
714 ///
715 /// No AD group, residual or gradient will share it, so it needs no
716 /// allocation group: the tensor's own storage returns to its pool when the
717 /// record drops, and a unique handle hands the tensor back unchanged.
718 Owned { tensor: Tensor },
719}
720
721/// Read-only descriptor record used by eager handles and the AD registries.
722#[derive(Debug)]
723pub(crate) struct AdValueRecord {
724 container: Arc<RetentionContainer>,
725 dtype: DType,
726 shape: Box<[usize]>,
727}
728
729impl AdValueRecord {
730 fn from_group(
731 group: AllocationGroup,
732 slot: DescriptorSlot,
733 dtype: DType,
734 shape: Vec<usize>,
735 ) -> Arc<Self> {
736 Arc::new(Self {
737 container: Arc::new(RetentionContainer::Pooled { group, slot }),
738 dtype,
739 shape: shape.into_boxed_slice(),
740 })
741 }
742
743 fn from_tensor(tensor: Tensor, op: &'static str) -> Result<Arc<Self>> {
744 let dtype = tensor.dtype();
745 let shape = tensor.shape().to_vec();
746 if matches!(dtype, DType::External(_)) {
747 // A caller-owned payload is retained directly: it owns no pooled
748 // storage, so there is no group to build and nothing to return to a
749 // pool when the record drops.
750 return Ok(Arc::new(Self {
751 container: Arc::new(RetentionContainer::CallerOwned {
752 tensor: Box::new(tensor),
753 }),
754 dtype,
755 shape: shape.into_boxed_slice(),
756 }));
757 }
758 let (group, bindings) = AllocationGroup::from_tensors(vec![tensor])
759 .map_err(|error| Error::runtime_state_source(op, ErrorPhase::Execution, error))?;
760 let slot = bindings.first().copied().ok_or_else(|| {
761 Error::runtime_state(op, ErrorPhase::Execution, "empty allocation-group binding")
762 })?;
763 Ok(Self::from_group(group, slot, dtype, shape))
764 }
765
766 /// Retain an untracked result without building an allocation group.
767 fn from_untracked_tensor(tensor: Tensor, op: &'static str) -> Result<Arc<Self>> {
768 if matches!(tensor.dtype(), DType::External(_)) {
769 return Self::from_tensor(tensor, op);
770 }
771 let dtype = tensor.dtype();
772 let shape = tensor.shape().to_vec().into_boxed_slice();
773 Ok(Arc::new(Self {
774 container: Arc::new(RetentionContainer::Owned { tensor }),
775 dtype,
776 shape,
777 }))
778 }
779
780 fn tensor_read(&self, op: &'static str) -> Result<TensorRead<'_>> {
781 match self.container.as_ref() {
782 RetentionContainer::Pooled { group, slot } => group.read_view(*slot).map_err(|error| {
783 // One slot, so no `Vec` and no empty-result case: the single-slot
784 // accessor is the whole request (#2079).
785 Error::runtime_state_source(op, ErrorPhase::Execution, error)
786 }),
787 RetentionContainer::CallerOwned { tensor } => Ok(TensorRead::from_tensor(tensor)),
788 RetentionContainer::Owned { tensor } => Ok(TensorRead::from_tensor(tensor)),
789 }
790 }
791
792 fn value(&self, op: &'static str) -> Result<ValueGuard<'_>> {
793 if let RetentionContainer::Owned { tensor } = self.container.as_ref() {
794 // A preset-dtype owned tensor always has a typed borrowed view.
795 return Ok(ValueGuard {
796 view: TensorRead::from_tensor(tensor).tensor_view(),
797 });
798 }
799 match self.tensor_read(op)? {
800 TensorRead::View(view) => Ok(ValueGuard { view }),
801 // A caller-owned payload has no typed descriptor view, so a path that
802 // needs one fails explicitly instead of borrowing the payload as bytes.
803 TensorRead::Tensor(_) => Err(Error::runtime_state(
804 op,
805 ErrorPhase::Execution,
806 "allocation-group value did not produce a borrowed descriptor view",
807 )),
808 }
809 }
810
811 fn dtype(&self) -> DType {
812 self.dtype
813 }
814
815 fn shape(&self) -> &[usize] {
816 &self.shape
817 }
818}
819
820/// Shared eager execution context for tensors on a backend.
821///
822/// Reusing one context lets eager tensors share backend state, extension
823/// runtime caches, and gradient storage across a computation.
824///
825/// # Examples
826///
827/// ```
828/// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
829///
830/// let ctx = EagerRuntime::new()?;
831/// let x = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![1], vec![1.0_f64]).unwrap(), ctx.clone()).unwrap();
832/// let y = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![1], vec![2.0_f64]).unwrap(), ctx.clone()).unwrap();
833/// let z = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |session| session.add(&x, &y)).unwrap();
834///
835/// assert_eq!(z.value().unwrap().as_slice::<f64>().unwrap(), &[3.0]);
836/// # Ok::<(), tenferro_ad::Error>(())
837/// ```
838pub struct EagerRuntime {
839 id: ContextId,
840 /// Prepared-program runtime. Engine kinds are registered on first use;
841 /// the runtime never stores an engine.
842 runtime: Runtime,
843 /// Serializes runtime reconfiguration (engine-kind registration and
844 /// extension-module installation).
845 extension_install_lock: Mutex<()>,
846 semantic_extension_rules: SemanticExtensionRuleSet,
847 grad_slots: Mutex<HashMap<ValueKey<StdTensorOp>, WeakGradSlot>>,
848 value_records: Mutex<HashMap<ValueKey<StdTensorOp>, Weak<EagerTensorRecord>>>,
849 ad_transform_cache: Arc<AdTransformCache>,
850 /// S2: prepared derivative programs keyed by semantic structure, wrt input,
851 /// and concrete bound input metadata. Avoids re-running freeze+AD
852 /// transform+compile_frozen on warm structure hits.
853 prepared_derivative_cache: Mutex<PreparedDerivativeCache>,
854}
855
856/// An eager runtime and its borrowed backend session for one execution boundary.
857///
858/// Obtain this only through [`EagerRuntime::with_eager_session`]. It rejects
859/// tensors from another eager runtime even if both runtimes use the same backend
860/// type, and it cannot escape the boundary closure.
861///
862/// # Examples
863///
864/// ```rust
865/// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
866///
867/// let ctx = EagerRuntime::new()?;
868/// let x = EagerTensor::from_tensor_in(
869/// Tensor::from_vec_col_major(vec![1], vec![3.0_f64])?, ctx.clone(),
870/// )?;
871/// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |session| session.neg(&x))?;
872/// assert_eq!(y.value()?.as_slice::<f64>()?, &[-3.0]);
873/// # Ok::<(), tenferro_ad::Error>(())
874/// ```
875///
876/// The borrowed session cannot escape its execution boundary:
877///
878/// ```compile_fail
879/// use tenferro_ad::EagerRuntime;
880/// let ctx = EagerRuntime::new().unwrap();
881/// let escaped = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |session| session).unwrap();
882/// let _ = escaped;
883/// ```
884pub struct EagerSession<'a> {
885 runtime: &'a Arc<EagerRuntime>,
886 engine: &'a mut dyn ExecutionEngine,
887}
888
889impl fmt::Debug for EagerSession<'_> {
890 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
891 f.debug_struct("EagerSession")
892 .field("runtime_id", &self.runtime.id())
893 .finish_non_exhaustive()
894 }
895}
896
897impl EagerSession<'_> {
898 /// Negate an eager tensor inside the caller's execution boundary.
899 ///
900 /// # Examples
901 ///
902 /// ```rust
903 /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
904 ///
905 /// let ctx = EagerRuntime::new()?;
906 /// let x = EagerTensor::from_tensor_in(
907 /// Tensor::from_vec_col_major(vec![1], vec![4.0_f64])?, ctx.clone(),
908 /// )?;
909 /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |session| session.neg(&x))?;
910 /// assert_eq!(y.value()?.as_slice::<f64>()?, &[-4.0]);
911 /// # Ok::<(), tenferro_ad::Error>(())
912 /// ```
913 ///
914 /// # Errors
915 ///
916 /// Returns [`Error::ContextMismatch`] for a tensor from another runtime,
917 /// or a typed eager/backend error from the selected operation.
918 pub fn neg(&mut self, input: &EagerTensor) -> Result<EagerTensor> {
919 self.run_unary(input, StdTensorOp::Neg)
920 }
921
922 /// Compute the elementwise exponential inside this eager session.
923 ///
924 /// # Examples
925 ///
926 /// ```rust
927 /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
928 /// let ctx = EagerRuntime::new()?;
929 /// let x = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![1], vec![0.0_f64])?, ctx.clone())?;
930 /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |session| session.exp(&x))?;
931 /// assert_eq!(y.value()?.as_slice::<f64>()?, &[1.0]);
932 /// # Ok::<(), tenferro_ad::Error>(())
933 /// ```
934 ///
935 /// # Errors
936 ///
937 /// Returns [`Error::ContextMismatch`] for a foreign runtime, or a typed
938 /// unsupported/backend error for the input dtype.
939 pub fn exp(&mut self, input: &EagerTensor) -> Result<EagerTensor> {
940 self.run_unary(input, StdTensorOp::Exp)
941 }
942
943 /// Compute the elementwise absolute value inside this eager session.
944 ///
945 /// # Examples
946 ///
947 /// ```rust
948 /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
949 /// let ctx = EagerRuntime::new()?;
950 /// let x = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![1], vec![-2.0_f64])?, ctx.clone())?;
951 /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |session| session.abs(&x))?;
952 /// assert_eq!(y.value()?.as_slice::<f64>()?, &[2.0]);
953 /// # Ok::<(), tenferro_ad::Error>(())
954 /// ```
955 ///
956 /// # Errors
957 ///
958 /// Returns [`Error::ContextMismatch`] for a foreign runtime, or a typed
959 /// unsupported/backend error for the input dtype.
960 pub fn abs(&mut self, input: &EagerTensor) -> Result<EagerTensor> {
961 self.run_unary(input, StdTensorOp::Abs)
962 }
963
964 /// Compute the elementwise conjugate inside this eager session.
965 ///
966 /// # Examples
967 ///
968 /// ```rust
969 /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
970 /// let ctx = EagerRuntime::new()?;
971 /// let x = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![1], vec![2.0_f64])?, ctx.clone())?;
972 /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |session| session.conj(&x))?;
973 /// assert_eq!(y.value()?.as_slice::<f64>()?, &[2.0]);
974 /// # Ok::<(), tenferro_ad::Error>(())
975 /// ```
976 ///
977 /// # Errors
978 ///
979 /// Returns [`Error::ContextMismatch`] for a foreign runtime, or a typed
980 /// unsupported/backend error for the input dtype.
981 pub fn conj(&mut self, input: &EagerTensor) -> Result<EagerTensor> {
982 self.run_unary(input, StdTensorOp::Conj)
983 }
984
985 /// Compute the elementwise sign on this borrowed session.
986 ///
987 /// # Examples
988 /// ```rust
989 /// use tenferro_ad::{EagerRuntime, Tensor};
990 /// let ctx = EagerRuntime::new()?;
991 /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
992 /// let x = s.constant_from(Tensor::from_vec_col_major(vec![1], vec![-2.0_f64])?)?;
993 /// s.sign(&x)
994 /// })?;
995 /// assert_eq!(y.value()?.as_slice::<f64>()?, &[-1.0]);
996 /// # Ok::<(), tenferro_ad::Error>(())
997 /// ```
998 /// # Errors
999 /// Returns a typed foreign-runtime, unsupported-dtype, or backend error.
1000 pub fn sign(&mut self, input: &EagerTensor) -> Result<EagerTensor> {
1001 self.run_unary(input, StdTensorOp::Sign)
1002 }
1003
1004 /// Compute the elementwise natural logarithm on this borrowed session.
1005 ///
1006 /// # Examples
1007 /// ```rust
1008 /// use tenferro_ad::{EagerRuntime, Tensor};
1009 /// let ctx = EagerRuntime::new()?;
1010 /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
1011 /// let x = s.constant_from(Tensor::from_vec_col_major(vec![1], vec![1.0_f64])?)?;
1012 /// s.log(&x)
1013 /// })?;
1014 /// assert_eq!(y.value()?.as_slice::<f64>()?, &[0.0]);
1015 /// # Ok::<(), tenferro_ad::Error>(())
1016 /// ```
1017 /// # Errors
1018 /// Returns a typed foreign-runtime, unsupported-dtype, or backend error.
1019 pub fn log(&mut self, input: &EagerTensor) -> Result<EagerTensor> {
1020 self.run_unary(input, StdTensorOp::Log)
1021 }
1022
1023 /// Compute the elementwise square root on this borrowed session.
1024 ///
1025 /// # Examples
1026 /// ```rust
1027 /// use tenferro_ad::{EagerRuntime, Tensor};
1028 /// let ctx = EagerRuntime::new()?;
1029 /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
1030 /// let x = s.constant_from(Tensor::from_vec_col_major(vec![1], vec![4.0_f64])?)?;
1031 /// s.sqrt(&x)
1032 /// })?;
1033 /// assert_eq!(y.value()?.as_slice::<f64>()?, &[2.0]);
1034 /// # Ok::<(), tenferro_ad::Error>(())
1035 /// ```
1036 /// # Errors
1037 /// Returns a typed foreign-runtime, unsupported-dtype, or backend error.
1038 pub fn sqrt(&mut self, input: &EagerTensor) -> Result<EagerTensor> {
1039 self.run_unary(input, StdTensorOp::Sqrt)
1040 }
1041
1042 /// Compute the elementwise reciprocal square root on this borrowed session.
1043 ///
1044 /// # Examples
1045 /// ```rust
1046 /// use tenferro_ad::{EagerRuntime, Tensor};
1047 /// let ctx = EagerRuntime::new()?;
1048 /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
1049 /// let x = s.constant_from(Tensor::from_vec_col_major(vec![1], vec![4.0_f64])?)?;
1050 /// s.rsqrt(&x)
1051 /// })?;
1052 /// assert_eq!(y.value()?.as_slice::<f64>()?, &[0.5]);
1053 /// # Ok::<(), tenferro_ad::Error>(())
1054 /// ```
1055 /// # Errors
1056 /// Returns a typed foreign-runtime, unsupported-dtype, or backend error.
1057 pub fn rsqrt(&mut self, input: &EagerTensor) -> Result<EagerTensor> {
1058 self.run_unary(input, StdTensorOp::Rsqrt)
1059 }
1060
1061 /// Compute the elementwise sine on this borrowed session.
1062 ///
1063 /// # Examples
1064 /// ```rust
1065 /// use tenferro_ad::{EagerRuntime, Tensor};
1066 /// let ctx = EagerRuntime::new()?;
1067 /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
1068 /// let x = s.constant_from(Tensor::from_vec_col_major(vec![1], vec![0.0_f64])?)?;
1069 /// s.sin(&x)
1070 /// })?;
1071 /// assert_eq!(y.value()?.as_slice::<f64>()?, &[0.0]);
1072 /// # Ok::<(), tenferro_ad::Error>(())
1073 /// ```
1074 /// # Errors
1075 /// Returns a typed foreign-runtime, unsupported-dtype, or backend error.
1076 pub fn sin(&mut self, input: &EagerTensor) -> Result<EagerTensor> {
1077 self.run_unary(input, StdTensorOp::Sin)
1078 }
1079
1080 /// Compute the elementwise cosine on this borrowed session.
1081 ///
1082 /// # Examples
1083 /// ```rust
1084 /// use tenferro_ad::{EagerRuntime, Tensor};
1085 /// let ctx = EagerRuntime::new()?;
1086 /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
1087 /// let x = s.constant_from(Tensor::from_vec_col_major(vec![1], vec![0.0_f64])?)?;
1088 /// s.cos(&x)
1089 /// })?;
1090 /// assert_eq!(y.value()?.as_slice::<f64>()?, &[1.0]);
1091 /// # Ok::<(), tenferro_ad::Error>(())
1092 /// ```
1093 /// # Errors
1094 /// Returns a typed foreign-runtime, unsupported-dtype, or backend error.
1095 pub fn cos(&mut self, input: &EagerTensor) -> Result<EagerTensor> {
1096 self.run_unary(input, StdTensorOp::Cos)
1097 }
1098
1099 /// Compute the elementwise hyperbolic tangent on this borrowed session.
1100 ///
1101 /// # Examples
1102 /// ```rust
1103 /// use tenferro_ad::{EagerRuntime, Tensor};
1104 /// let ctx = EagerRuntime::new()?;
1105 /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
1106 /// let x = s.constant_from(Tensor::from_vec_col_major(vec![1], vec![0.0_f64])?)?;
1107 /// s.tanh(&x)
1108 /// })?;
1109 /// assert_eq!(y.value()?.as_slice::<f64>()?, &[0.0]);
1110 /// # Ok::<(), tenferro_ad::Error>(())
1111 /// ```
1112 /// # Errors
1113 /// Returns a typed foreign-runtime, unsupported-dtype, or backend error.
1114 pub fn tanh(&mut self, input: &EagerTensor) -> Result<EagerTensor> {
1115 self.run_unary(input, StdTensorOp::Tanh)
1116 }
1117
1118 /// Compute `exp(x) - 1` elementwise on this borrowed session.
1119 ///
1120 /// # Examples
1121 /// ```rust
1122 /// use tenferro_ad::{EagerRuntime, Tensor};
1123 /// let ctx = EagerRuntime::new()?;
1124 /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
1125 /// let x = s.constant_from(Tensor::from_vec_col_major(vec![1], vec![0.0_f64])?)?;
1126 /// s.expm1(&x)
1127 /// })?;
1128 /// assert_eq!(y.value()?.as_slice::<f64>()?, &[0.0]);
1129 /// # Ok::<(), tenferro_ad::Error>(())
1130 /// ```
1131 /// # Errors
1132 /// Returns a typed foreign-runtime, unsupported-dtype, or backend error.
1133 pub fn expm1(&mut self, input: &EagerTensor) -> Result<EagerTensor> {
1134 self.run_unary(input, StdTensorOp::Expm1)
1135 }
1136
1137 /// Compute `log(1 + x)` elementwise on this borrowed session.
1138 ///
1139 /// # Examples
1140 /// ```rust
1141 /// use tenferro_ad::{EagerRuntime, Tensor};
1142 /// let ctx = EagerRuntime::new()?;
1143 /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
1144 /// let x = s.constant_from(Tensor::from_vec_col_major(vec![1], vec![0.0_f64])?)?;
1145 /// s.log1p(&x)
1146 /// })?;
1147 /// assert_eq!(y.value()?.as_slice::<f64>()?, &[0.0]);
1148 /// # Ok::<(), tenferro_ad::Error>(())
1149 /// ```
1150 /// # Errors
1151 /// Returns a typed foreign-runtime, unsupported-dtype, or backend error.
1152 pub fn log1p(&mut self, input: &EagerTensor) -> Result<EagerTensor> {
1153 self.run_unary(input, StdTensorOp::Log1p)
1154 }
1155
1156 /// Compute the error function `erf(x)` elementwise on this borrowed session.
1157 ///
1158 /// Defined for real `F32`/`F64` tensors; `erf(+-0) = +-0`,
1159 /// `erf(+-inf) = +-1`, and `NaN` stays `NaN`. The derivative is
1160 /// `2/sqrt(pi) * exp(-x^2)`.
1161 ///
1162 /// # Examples
1163 /// ```rust
1164 /// use tenferro_ad::{EagerRuntime, Tensor};
1165 /// let ctx = EagerRuntime::new()?;
1166 /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
1167 /// let x = s.constant_from(Tensor::from_vec_col_major(vec![2], vec![0.0_f64, 1.0])?)?;
1168 /// s.erf(&x)
1169 /// })?;
1170 /// let y = y.value()?;
1171 /// let y = y.as_slice::<f64>()?;
1172 /// assert_eq!(y[0], 0.0);
1173 /// assert!((y[1] - 0.842_700_792_949_714_9).abs() < 1.0e-15);
1174 /// # Ok::<(), tenferro_ad::Error>(())
1175 /// ```
1176 /// # Errors
1177 /// Returns a typed foreign-runtime error, a typed unsupported-dtype error
1178 /// for complex, integer, or `Bool` input, or a backend error.
1179 pub fn erf(&mut self, input: &EagerTensor) -> Result<EagerTensor> {
1180 self.run_unary(input, StdTensorOp::Erf)
1181 }
1182
1183 /// Convert a tensor under the checked dtype-promotion lattice.
1184 /// Use [`Self::cast`] for intentional lossy projection.
1185 ///
1186 /// # Examples
1187 ///
1188 /// ```rust
1189 /// use tenferro_ad::{DType, EagerRuntime, Tensor};
1190 /// let ctx = EagerRuntime::new()?;
1191 /// let converted = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |session| {
1192 /// let x = session.constant_from(Tensor::from_vec_col_major(vec![1], vec![2.0_f64])?)?;
1193 /// session.convert(&x, DType::C64)
1194 /// })?;
1195 /// assert_eq!(converted.dtype(), DType::C64);
1196 /// # Ok::<(), tenferro_ad::Error>(())
1197 /// ```
1198 ///
1199 /// # Errors
1200 ///
1201 /// Returns [`Error::ContextMismatch`] for a foreign runtime or a typed
1202 /// unsupported dtype conversion/backend error.
1203 pub fn convert(&mut self, input: &EagerTensor, to: DType) -> Result<EagerTensor> {
1204 self.ensure_runtime(input)?;
1205 tenferro_tensor::validate::validate_convert_dtype(
1206 "EagerTensor::convert",
1207 input.dtype(),
1208 to,
1209 )
1210 .map_err(Error::TensorRuntime)?;
1211 self.cast(input, to)
1212 }
1213
1214 /// Cast a tensor to a dtype, permitting explicitly lossy projections.
1215 ///
1216 /// # Examples
1217 ///
1218 /// ```rust
1219 /// use tenferro_ad::{DType, EagerRuntime, Tensor};
1220 /// let ctx = EagerRuntime::new()?;
1221 /// let casted = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |session| {
1222 /// let x = session.constant_from(Tensor::from_vec_col_major(vec![1], vec![2.8_f64])?)?;
1223 /// session.cast(&x, DType::I32)
1224 /// })?;
1225 /// assert_eq!(casted.value()?.as_slice::<i32>()?, &[2]);
1226 /// # Ok::<(), tenferro_ad::Error>(())
1227 /// ```
1228 ///
1229 /// # Errors
1230 ///
1231 /// Returns [`Error::ContextMismatch`] for a foreign runtime, or a typed
1232 /// unsupported projection/backend error.
1233 pub fn cast(&mut self, input: &EagerTensor, to: DType) -> Result<EagerTensor> {
1234 self.run_unary(
1235 input,
1236 StdTensorOp::Convert {
1237 from: input.dtype(),
1238 to,
1239 },
1240 )
1241 }
1242
1243 /// Permute the axes of an eager tensor while preserving independent ownership.
1244 ///
1245 /// # Examples
1246 ///
1247 /// ```rust
1248 /// use tenferro_ad::{EagerRuntime, Tensor};
1249 /// let ctx = EagerRuntime::new()?;
1250 /// let copied = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |session| {
1251 /// let x = session.constant_from(Tensor::from_vec_col_major(vec![2, 2], vec![1.0_f64, 2.0, 3.0, 4.0])?)?;
1252 /// let y = session.transpose(&x, &[1, 0])?;
1253 /// session.duplicate_value(&y)
1254 /// })?;
1255 /// assert_eq!(copied.as_slice::<f64>()?, &[1.0, 3.0, 2.0, 4.0]);
1256 /// # Ok::<(), tenferro_ad::Error>(())
1257 /// ```
1258 ///
1259 /// # Errors
1260 ///
1261 /// Returns [`Error::ContextMismatch`] for a foreign runtime, or a typed
1262 /// axis/backend error for an invalid permutation or copy.
1263 pub fn transpose(&mut self, input: &EagerTensor, perm: &[usize]) -> Result<EagerTensor> {
1264 self.ensure_runtime(input)?;
1265 input.transpose_in_session(perm, &mut *self.engine)
1266 }
1267
1268 /// Reshape an eager tensor while retaining a separate result owner.
1269 ///
1270 /// # Examples
1271 ///
1272 /// ```rust
1273 /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
1274 /// let ctx = EagerRuntime::new()?;
1275 /// let x = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0])?, ctx.clone())?;
1276 /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |session| session.reshape(&x, [1, 2]))?;
1277 /// assert_eq!(y.shape(), &[1, 2]);
1278 /// assert_eq!(y.value()?.as_slice::<f64>()?, &[1.0, 2.0]);
1279 /// # Ok::<(), tenferro_ad::Error>(())
1280 /// ```
1281 ///
1282 /// # Errors
1283 ///
1284 /// Returns [`Error::ContextMismatch`] for a foreign runtime, or a typed
1285 /// validation/backend error when the target shape is incompatible.
1286 pub fn reshape(
1287 &mut self,
1288 input: &EagerTensor,
1289 shape: impl IntoShapeVec,
1290 ) -> Result<EagerTensor> {
1291 self.ensure_runtime(input)?;
1292 input.reshape_in_session(&shape.into_shape_vec(), &mut *self.engine)
1293 }
1294
1295 /// Slice an eager tensor with explicit start, limit, and stride per axis.
1296 ///
1297 /// # Examples
1298 ///
1299 /// ```rust
1300 /// use tenferro_ad::{EagerRuntime, SliceConfig, Tensor};
1301 /// let ctx = EagerRuntime::new()?;
1302 /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |session| {
1303 /// let x = session.constant_from(Tensor::from_vec_col_major(vec![4], vec![1.0_f64, 2.0, 3.0, 4.0])?)?;
1304 /// session.slice(&x, SliceConfig { starts: vec![1], limits: vec![3], strides: vec![1] })
1305 /// })?;
1306 /// assert_eq!(y.value()?.as_slice::<f64>()?, &[2.0, 3.0]);
1307 /// # Ok::<(), tenferro_ad::Error>(())
1308 /// ```
1309 ///
1310 /// # Errors
1311 ///
1312 /// Returns [`Error::ContextMismatch`] for a foreign runtime, or a typed
1313 /// axis/stride/backend error for an invalid slice or copy.
1314 pub fn slice(&mut self, input: &EagerTensor, config: SliceConfig) -> Result<EagerTensor> {
1315 self.ensure_runtime(input)?;
1316 input.slice_in_session(config, &mut *self.engine)
1317 }
1318
1319 /// Broadcast an eager tensor into a larger shape on this session.
1320 ///
1321 /// # Examples
1322 ///
1323 /// ```rust
1324 /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
1325 /// let ctx = EagerRuntime::new()?;
1326 /// let x = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0])?, ctx.clone())?;
1327 /// let copy = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |session| {
1328 /// let y = session.broadcast_in_dim(&x, &[2, 2], &[0])?;
1329 /// session.duplicate_value(&y)
1330 /// })?;
1331 /// assert_eq!(copy.as_slice::<f64>()?, &[1.0, 2.0, 1.0, 2.0]);
1332 /// # Ok::<(), tenferro_ad::Error>(())
1333 /// ```
1334 ///
1335 /// # Errors
1336 ///
1337 /// Returns [`Error::ContextMismatch`] for a foreign runtime, or a typed
1338 /// validation/backend error for an invalid broadcast mapping.
1339 pub fn broadcast_in_dim(
1340 &mut self,
1341 input: &EagerTensor,
1342 shape: &[usize],
1343 dims: &[usize],
1344 ) -> Result<EagerTensor> {
1345 self.ensure_runtime(input)?;
1346 input.broadcast_in_dim_in_session(shape, dims, &mut *self.engine)
1347 }
1348
1349 /// Keep the lower triangle of a matrix on this borrowed session.
1350 ///
1351 /// # Examples
1352 /// ```rust
1353 /// use tenferro_ad::{EagerRuntime, Tensor};
1354 /// let ctx = EagerRuntime::new()?;
1355 /// let lower = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
1356 /// let matrix = s.constant_from(Tensor::from_vec_col_major(vec![2, 2], vec![1.0_f64, 2.0, 3.0, 4.0])?)?;
1357 /// s.tril(&matrix, 0)
1358 /// })?;
1359 /// assert_eq!(lower.value()?.as_slice::<f64>()?, &[1.0, 2.0, 0.0, 4.0]);
1360 /// # Ok::<(), tenferro_ad::Error>(())
1361 /// ```
1362 /// # Errors
1363 /// Returns a typed foreign-runtime, rank, unsupported-dtype, or backend error.
1364 pub fn tril(&mut self, input: &EagerTensor, k: i64) -> Result<EagerTensor> {
1365 self.run_unary(input, StdTensorOp::Tril { k })
1366 }
1367
1368 /// Keep the upper triangle of a matrix on this borrowed session.
1369 ///
1370 /// # Examples
1371 /// ```rust
1372 /// use tenferro_ad::{EagerRuntime, Tensor};
1373 /// let ctx = EagerRuntime::new()?;
1374 /// let upper = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
1375 /// let matrix = s.constant_from(Tensor::from_vec_col_major(vec![2, 2], vec![1.0_f64, 2.0, 3.0, 4.0])?)?;
1376 /// s.triu(&matrix, 0)
1377 /// })?;
1378 /// assert_eq!(upper.value()?.as_slice::<f64>()?, &[1.0, 0.0, 3.0, 4.0]);
1379 /// # Ok::<(), tenferro_ad::Error>(())
1380 /// ```
1381 /// # Errors
1382 /// Returns a typed foreign-runtime, rank, unsupported-dtype, or backend error.
1383 pub fn triu(&mut self, input: &EagerTensor, k: i64) -> Result<EagerTensor> {
1384 self.run_unary(input, StdTensorOp::Triu { k })
1385 }
1386
1387 /// Pad an eager tensor with zeros on this borrowed session.
1388 ///
1389 /// # Examples
1390 /// ```rust
1391 /// use tenferro_ad::{EagerRuntime, PadConfig, Tensor};
1392 /// let ctx = EagerRuntime::new()?;
1393 /// let padded = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
1394 /// let x = s.constant_from(Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0])?)?;
1395 /// s.pad(&x, PadConfig {
1396 /// edge_padding_low: vec![1],
1397 /// edge_padding_high: vec![1],
1398 /// interior_padding: vec![1],
1399 /// })
1400 /// })?;
1401 /// assert_eq!(padded.value()?.as_slice::<f64>()?, &[0.0, 1.0, 0.0, 2.0, 0.0]);
1402 /// # Ok::<(), tenferro_ad::Error>(())
1403 /// ```
1404 /// # Errors
1405 /// Returns [`Error::ContextMismatch`] when an input belongs to another eager
1406 /// runtime, a validation error with
1407 /// `ValidationError::InvalidArgument` for a padding configuration whose
1408 /// length or extents do not match the input rank, or
1409 /// [`Error::TensorRuntime`] for a typed backend failure.
1410 pub fn pad(&mut self, input: &EagerTensor, config: PadConfig) -> Result<EagerTensor> {
1411 self.run_unary(input, StdTensorOp::Pad(config))
1412 }
1413
1414 /// Reverse the elements along selected axes on this borrowed session.
1415 ///
1416 /// # Examples
1417 /// ```rust
1418 /// use tenferro_ad::{EagerRuntime, Tensor};
1419 /// let ctx = EagerRuntime::new()?;
1420 /// let reversed = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
1421 /// let x = s.constant_from(Tensor::from_vec_col_major(vec![3], vec![1.0_f64, 2.0, 3.0])?)?;
1422 /// s.reverse(&x, &[0])
1423 /// })?;
1424 /// assert_eq!(reversed.value()?.as_slice::<f64>()?, &[3.0, 2.0, 1.0]);
1425 /// # Ok::<(), tenferro_ad::Error>(())
1426 /// ```
1427 /// # Errors
1428 /// Returns a typed foreign-runtime, invalid-axis, or backend error.
1429 pub fn reverse(&mut self, input: &EagerTensor, axes: &[usize]) -> Result<EagerTensor> {
1430 self.ensure_runtime(input)?;
1431 crate::eager_ops::validate_eager_axes("EagerSession::reverse", input.shape().len(), axes)?;
1432 self.run_unary(
1433 input,
1434 StdTensorOp::Reverse {
1435 axes: axes.to_vec(),
1436 },
1437 )
1438 }
1439
1440 /// Slice an eager tensor using runtime start indices in this borrowed session.
1441 ///
1442 /// # Examples
1443 /// ```rust
1444 /// use tenferro_ad::{EagerRuntime, Tensor};
1445 /// let ctx = EagerRuntime::new()?;
1446 /// let selected = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
1447 /// let x = s.constant_from(Tensor::from_vec_col_major(vec![4], vec![1.0_f64, 2.0, 3.0, 4.0])?)?;
1448 /// let starts = s.constant_from(Tensor::from_vec_col_major(vec![1], vec![1_i64])?)?;
1449 /// s.dynamic_slice(&x, &starts, &[2])
1450 /// })?;
1451 /// assert_eq!(selected.value()?.as_slice::<f64>()?, &[2.0, 3.0]);
1452 /// # Ok::<(), tenferro_ad::Error>(())
1453 /// ```
1454 /// # Errors
1455 /// Returns a typed foreign-runtime, invalid-index or slice-shape, or backend error.
1456 pub fn dynamic_slice(
1457 &mut self,
1458 input: &EagerTensor,
1459 starts: &EagerTensor,
1460 sizes: &[usize],
1461 ) -> Result<EagerTensor> {
1462 self.ensure_runtime(input)?;
1463 self.ensure_runtime(starts)?;
1464 EagerTensor::nary_op_in_session(
1465 &[input, starts],
1466 StdTensorOp::DynamicSlice {
1467 slice_sizes: sizes.to_vec(),
1468 },
1469 &mut *self.engine,
1470 )
1471 }
1472
1473 /// Gather elements of an eager tensor in this borrowed session.
1474 ///
1475 /// # Examples
1476 /// ```rust
1477 /// use tenferro_ad::{EagerRuntime, GatherConfig, Tensor};
1478 /// let ctx = EagerRuntime::new()?;
1479 /// let result = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
1480 /// let x = s.constant_from(Tensor::from_vec_col_major(vec![3], vec![10.0_f64, 20.0, 30.0])?)?;
1481 /// let indices = s.constant_from(Tensor::from_vec_col_major(vec![2, 1], vec![2_i64, 0])?)?;
1482 /// s.gather(&x, &indices, GatherConfig {
1483 /// offset_dims: vec![], collapsed_slice_dims: vec![0],
1484 /// start_index_map: vec![0], index_vector_dim: 1,
1485 /// slice_sizes: vec![1],
1486 /// })
1487 /// })?;
1488 /// assert_eq!(result.value()?.as_slice::<f64>()?, &[30.0, 10.0]);
1489 /// # Ok::<(), tenferro_ad::Error>(())
1490 /// ```
1491 /// # Errors
1492 /// Returns a typed foreign-runtime, invalid-index/configuration, or backend error.
1493 pub fn gather(
1494 &mut self,
1495 input: &EagerTensor,
1496 indices: &EagerTensor,
1497 config: GatherConfig,
1498 ) -> Result<EagerTensor> {
1499 self.ensure_runtime(input)?;
1500 self.ensure_runtime(indices)?;
1501 EagerTensor::nary_op_in_session(
1502 &[input, indices],
1503 StdTensorOp::Gather(config),
1504 &mut *self.engine,
1505 )
1506 }
1507
1508 /// Concatenate eager tensors along one axis in this borrowed session.
1509 ///
1510 /// # Examples
1511 /// ```rust
1512 /// use tenferro_ad::{EagerRuntime, Tensor};
1513 /// let ctx = EagerRuntime::new()?;
1514 /// let result = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
1515 /// let a = s.constant_from(Tensor::from_vec_col_major(vec![1], vec![1.0_f64])?)?;
1516 /// let b = s.constant_from(Tensor::from_vec_col_major(vec![1], vec![2.0_f64])?)?;
1517 /// s.concatenate(&[&a, &b], 0)
1518 /// })?;
1519 /// assert_eq!(result.value()?.as_slice::<f64>()?, &[1.0, 2.0]);
1520 /// # Ok::<(), tenferro_ad::Error>(())
1521 /// ```
1522 /// # Errors
1523 /// Returns a typed empty-input, foreign-runtime, invalid-axis/shape, or backend error.
1524 pub fn concatenate(&mut self, inputs: &[&EagerTensor], axis: usize) -> Result<EagerTensor> {
1525 for input in inputs {
1526 self.ensure_runtime(input)?;
1527 }
1528 EagerTensor::nary_op_in_session(
1529 inputs,
1530 StdTensorOp::Concatenate {
1531 axis,
1532 input_count: inputs.len(),
1533 },
1534 &mut *self.engine,
1535 )
1536 }
1537
1538 /// Scatter updates into an eager tensor within this borrowed session.
1539 ///
1540 /// # Examples
1541 /// ```rust
1542 /// use tenferro_ad::{EagerRuntime, ScatterConfig, Tensor};
1543 /// let ctx = EagerRuntime::new()?;
1544 /// let result = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
1545 /// let input = s.constant_from(Tensor::from_vec_col_major(vec![4], vec![0.0_f64; 4])?)?;
1546 /// let indices = s.constant_from(Tensor::from_vec_col_major(vec![2, 1], vec![1_i64, 3])?)?;
1547 /// let updates = s.constant_from(Tensor::from_vec_col_major(vec![2], vec![5.0_f64, 7.0])?)?;
1548 /// s.scatter(&input, &indices, &updates, ScatterConfig {
1549 /// update_window_dims: vec![],
1550 /// inserted_window_dims: vec![0],
1551 /// scatter_dims_to_operand_dims: vec![0],
1552 /// index_vector_dim: 1,
1553 /// })
1554 /// })?;
1555 /// assert_eq!(result.value()?.as_slice::<f64>()?, &[0.0, 5.0, 0.0, 7.0]);
1556 /// # Ok::<(), tenferro_ad::Error>(())
1557 /// ```
1558 /// # Errors
1559 /// Returns a typed foreign-runtime, invalid-index/configuration, or backend error.
1560 pub fn scatter(
1561 &mut self,
1562 input: &EagerTensor,
1563 indices: &EagerTensor,
1564 updates: &EagerTensor,
1565 config: ScatterConfig,
1566 ) -> Result<EagerTensor> {
1567 self.ensure_runtime(input)?;
1568 self.ensure_runtime(indices)?;
1569 self.ensure_runtime(updates)?;
1570 EagerTensor::nary_op_in_session(
1571 &[input, indices, updates],
1572 StdTensorOp::Scatter(config),
1573 &mut *self.engine,
1574 )
1575 }
1576
1577 /// Extract a diagonal along two axes in this borrowed session.
1578 ///
1579 /// # Examples
1580 /// ```rust
1581 /// use tenferro_ad::{EagerRuntime, Tensor};
1582 /// let ctx = EagerRuntime::new()?;
1583 /// let diagonal = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
1584 /// let matrix = s.constant_from(Tensor::from_vec_col_major(vec![2, 2], vec![1.0_f64, 2.0, 3.0, 4.0])?)?;
1585 /// s.extract_diag(&matrix, 0, 1)
1586 /// })?;
1587 /// assert_eq!(diagonal.value()?.as_slice::<f64>()?, &[1.0, 4.0]);
1588 /// # Ok::<(), tenferro_ad::Error>(())
1589 /// ```
1590 /// # Errors
1591 /// Returns a typed foreign-runtime, invalid-axis, or backend error.
1592 pub fn extract_diag(
1593 &mut self,
1594 input: &EagerTensor,
1595 axis_a: usize,
1596 axis_b: usize,
1597 ) -> Result<EagerTensor> {
1598 self.ensure_runtime(input)?;
1599 self.run_unary(input, StdTensorOp::ExtractDiag { axis_a, axis_b })
1600 }
1601
1602 /// Embed the input along a diagonal in this borrowed session.
1603 ///
1604 /// # Examples
1605 /// ```rust
1606 /// use tenferro_ad::{EagerRuntime, Tensor};
1607 /// let ctx = EagerRuntime::new()?;
1608 /// let matrix = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
1609 /// let diagonal = s.constant_from(Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0])?)?;
1610 /// s.embed_diag(&diagonal, 0, 1)
1611 /// })?;
1612 /// assert_eq!(matrix.value()?.as_slice::<f64>()?, &[1.0, 0.0, 0.0, 2.0]);
1613 /// # Ok::<(), tenferro_ad::Error>(())
1614 /// ```
1615 /// # Errors
1616 /// Returns a typed foreign-runtime, invalid-axis, or backend error.
1617 pub fn embed_diag(
1618 &mut self,
1619 input: &EagerTensor,
1620 axis_a: usize,
1621 axis_b: usize,
1622 ) -> Result<EagerTensor> {
1623 self.ensure_runtime(input)?;
1624 self.run_unary(input, StdTensorOp::EmbedDiag { axis_a, axis_b })
1625 }
1626
1627 /// Reduce an eager tensor over selected axes within this borrowed session.
1628 /// `None` reduces all axes, while `Some(&[])` retains the input shape.
1629 ///
1630 /// # Examples
1631 ///
1632 /// ```rust
1633 /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
1634 /// let ctx = EagerRuntime::new()?;
1635 /// let x = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0])?, ctx.clone())?;
1636 /// let sum = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |session| session.reduce_sum(&x, None))?;
1637 /// assert_eq!(sum.value()?.as_slice::<f64>()?, &[3.0]);
1638 /// # Ok::<(), tenferro_ad::Error>(())
1639 /// ```
1640 ///
1641 /// # Errors
1642 ///
1643 /// Returns [`Error::ContextMismatch`] for a foreign runtime, or a typed
1644 /// validation/backend error for invalid axes or unsupported dtypes.
1645 pub fn reduce_sum(
1646 &mut self,
1647 input: &EagerTensor,
1648 axes: Option<&[usize]>,
1649 ) -> Result<EagerTensor> {
1650 self.ensure_runtime(input)?;
1651 input.reduce_sum_in_session(axes, &mut *self.engine)
1652 }
1653
1654 /// Sum elementwise squares over the selected axes in this borrowed session.
1655 /// Only `f32` and `f64` are supported. `None` reduces every axis, like the
1656 /// rest of the reduction family; `Some(&[])` squares each value.
1657 ///
1658 /// # Examples
1659 /// ```rust
1660 /// use tenferro_ad::{EagerRuntime, Tensor};
1661 /// let ctx = EagerRuntime::new()?;
1662 /// let (sum, all) = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
1663 /// let input = s.constant_from(Tensor::from_vec_col_major([2], vec![3.0_f64, 4.0])?)?;
1664 /// Ok::<_, tenferro_ad::Error>((
1665 /// s.reduce_sum_squares(&input, Some(&[0]))?,
1666 /// s.reduce_sum_squares(&input, None)?,
1667 /// ))
1668 /// })?;
1669 /// assert_eq!(sum.value()?.as_slice::<f64>()?, &[25.0]);
1670 /// assert_eq!(all.value()?.as_slice::<f64>()?, &[25.0]);
1671 /// # Ok::<(), tenferro_ad::Error>(())
1672 /// ```
1673 /// # Errors
1674 /// Returns typed foreign-runtime, invalid-axis, unsupported-dtype, or backend errors.
1675 pub fn reduce_sum_squares(
1676 &mut self,
1677 input: &EagerTensor,
1678 axes: Option<&[usize]>,
1679 ) -> Result<EagerTensor> {
1680 self.ensure_runtime(input)?;
1681 let axes = axes.map_or_else(|| (0..input.shape().len()).collect(), <[usize]>::to_vec);
1682 crate::eager_ops::validate_eager_axes(
1683 "EagerSession::reduce_sum_squares",
1684 input.shape().len(),
1685 &axes,
1686 )?;
1687 self.run_unary(input, StdTensorOp::ReduceSumSquares { axes })
1688 }
1689
1690 /// Reduce the product of selected axes in this borrowed session.
1691 /// `None` reduces every axis.
1692 ///
1693 /// # Examples
1694 /// ```rust
1695 /// use tenferro_ad::{EagerRuntime, Tensor};
1696 /// let ctx = EagerRuntime::new()?;
1697 /// let result = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
1698 /// let x = s.constant_from(Tensor::from_vec_col_major(vec![2], vec![2.0_f64, 3.0])?)?;
1699 /// s.reduce_prod(&x, None)
1700 /// })?;
1701 /// assert_eq!(result.value()?.as_slice::<f64>()?, &[6.0]);
1702 /// # Ok::<(), tenferro_ad::Error>(())
1703 /// ```
1704 /// # Errors
1705 /// Returns a typed foreign-runtime, invalid-axis, unsupported-dtype, or backend error.
1706 pub fn reduce_prod(
1707 &mut self,
1708 input: &EagerTensor,
1709 axes: Option<&[usize]>,
1710 ) -> Result<EagerTensor> {
1711 self.ensure_runtime(input)?;
1712 let axes = axes.map_or_else(|| (0..input.shape().len()).collect(), <[usize]>::to_vec);
1713 crate::eager_ops::validate_eager_axes(
1714 "EagerSession::reduce_prod",
1715 input.shape().len(),
1716 &axes,
1717 )?;
1718 self.run_unary(input, StdTensorOp::ReduceProd { axes })
1719 }
1720
1721 /// Reduce the maximum over selected axes in this borrowed session.
1722 /// `None` reduces every axis.
1723 ///
1724 /// # Examples
1725 /// ```rust
1726 /// use tenferro_ad::{EagerRuntime, Tensor};
1727 /// let ctx = EagerRuntime::new()?;
1728 /// let result = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
1729 /// let x = s.constant_from(Tensor::from_vec_col_major(vec![2], vec![2.0_f64, 3.0])?)?;
1730 /// s.reduce_max(&x, None)
1731 /// })?;
1732 /// assert_eq!(result.value()?.as_slice::<f64>()?, &[3.0]);
1733 /// # Ok::<(), tenferro_ad::Error>(())
1734 /// ```
1735 /// # Errors
1736 /// Returns a typed foreign-runtime, invalid-axis, unsupported-dtype, or backend error.
1737 pub fn reduce_max(
1738 &mut self,
1739 input: &EagerTensor,
1740 axes: Option<&[usize]>,
1741 ) -> Result<EagerTensor> {
1742 self.ensure_runtime(input)?;
1743 let axes = axes.map_or_else(|| (0..input.shape().len()).collect(), <[usize]>::to_vec);
1744 crate::eager_ops::validate_eager_axes(
1745 "EagerSession::reduce_max",
1746 input.shape().len(),
1747 &axes,
1748 )?;
1749 self.run_unary(input, StdTensorOp::ReduceMax { axes })
1750 }
1751
1752 /// Reduce the minimum over selected axes in this borrowed session.
1753 /// `None` reduces every axis.
1754 ///
1755 /// # Examples
1756 /// ```rust
1757 /// use tenferro_ad::{EagerRuntime, Tensor};
1758 /// let ctx = EagerRuntime::new()?;
1759 /// let result = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
1760 /// let x = s.constant_from(Tensor::from_vec_col_major(vec![2], vec![2.0_f64, 3.0])?)?;
1761 /// s.reduce_min(&x, None)
1762 /// })?;
1763 /// assert_eq!(result.value()?.as_slice::<f64>()?, &[2.0]);
1764 /// # Ok::<(), tenferro_ad::Error>(())
1765 /// ```
1766 /// # Errors
1767 /// Returns a typed foreign-runtime, invalid-axis, unsupported-dtype, or backend error.
1768 pub fn reduce_min(
1769 &mut self,
1770 input: &EagerTensor,
1771 axes: Option<&[usize]>,
1772 ) -> Result<EagerTensor> {
1773 self.ensure_runtime(input)?;
1774 let axes = axes.map_or_else(|| (0..input.shape().len()).collect(), <[usize]>::to_vec);
1775 crate::eager_ops::validate_eager_axes(
1776 "EagerSession::reduce_min",
1777 input.shape().len(),
1778 &axes,
1779 )?;
1780 self.run_unary(input, StdTensorOp::ReduceMin { axes })
1781 }
1782
1783 /// Duplicate an eager value into an independent tensor within the caller's
1784 /// borrowed session, preserving its dtype and placement.
1785 ///
1786 /// # Examples
1787 ///
1788 /// ```rust
1789 /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
1790 /// let ctx = EagerRuntime::new()?;
1791 /// let x = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![1], vec![2.0_f64])?, ctx.clone())?;
1792 /// let copy = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |session| session.duplicate_value(&x))?;
1793 /// assert_eq!(copy.as_slice::<f64>()?, &[2.0]);
1794 /// # Ok::<(), tenferro_ad::Error>(())
1795 /// ```
1796 ///
1797 /// # Errors
1798 ///
1799 /// Returns [`Error::ContextMismatch`] for a foreign runtime, or a typed
1800 /// runtime/backend error when the retained value cannot be duplicated.
1801 pub fn duplicate_value(&mut self, input: &EagerTensor) -> Result<Tensor> {
1802 self.ensure_runtime(input)?;
1803 input.duplicate_value_in_session(&mut *self.engine)
1804 }
1805
1806 /// Import an untracked leaf within this borrowed session.
1807 ///
1808 /// `tensor` must already be usable by this runtime's backend: a host
1809 /// tensor on a CPU runtime, or a tensor already on the device of a CUDA
1810 /// or WebGPU runtime. No host/device transfer happens here. To import host
1811 /// data into a device runtime, use [`Self::constant_from_host`], which
1812 /// uploads first; on a CPU runtime the two are equivalent.
1813 ///
1814 /// # Examples
1815 ///
1816 /// ```rust
1817 /// use tenferro_ad::{EagerRuntime, Tensor};
1818 /// let ctx = EagerRuntime::new()?;
1819 /// let c = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |session| {
1820 /// session.constant_from(Tensor::from_vec_col_major(vec![1], vec![2.0_f64])?)
1821 /// })?;
1822 /// assert_eq!(c.value()?.as_slice::<f64>()?, &[2.0]);
1823 /// # Ok::<(), tenferro_ad::Error>(())
1824 /// ```
1825 ///
1826 /// # Errors
1827 ///
1828 /// Returns [`Error::TensorRuntime`] for a typed backend failure when the value cannot be
1829 /// registered in the session, or [`Error::RuntimeState`] when the runtime's
1830 /// value registry is unavailable.
1831 pub fn constant_from(&mut self, tensor: Tensor) -> Result<EagerTensor> {
1832 EagerTensor::new_leaf_in_session(Arc::clone(self.runtime), tensor, false, &mut *self.engine)
1833 }
1834
1835 /// Upload a host tensor and import it as an untracked leaf in this session.
1836 ///
1837 /// Unlike [`Self::constant_from`], this explicitly crosses the host/device
1838 /// boundary: the host `tensor` is uploaded to this runtime's backend
1839 /// (a host copy on a CPU runtime) and the uploaded value becomes the leaf.
1840 /// Use it whenever the source data lives on the host and the runtime may
1841 /// be a device runtime; use [`Self::constant_from`] for a tensor that is
1842 /// already resident on the backend.
1843 ///
1844 /// # Examples
1845 /// ```rust
1846 /// use tenferro_ad::{EagerRuntime, Tensor};
1847 /// let ctx = EagerRuntime::new()?;
1848 /// let c = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
1849 /// s.constant_from_host(Tensor::from_vec_col_major([1], vec![2.0_f64])?)
1850 /// })?;
1851 /// assert_eq!(c.value()?.as_slice::<f64>()?, &[2.0]);
1852 /// # Ok::<(), tenferro_ad::Error>(())
1853 /// ```
1854 /// # Errors
1855 /// Returns [`Error::TensorRuntime`] for a typed backend failure, including a host-tensor
1856 /// upload failure, or [`Error::RuntimeState`] when the runtime's value
1857 /// registry is unavailable.
1858 pub fn constant_from_host(&mut self, tensor: Tensor) -> Result<EagerTensor> {
1859 let uploaded = self
1860 .engine
1861 .upload_host_tensor(TensorRead::from_tensor(&tensor))
1862 .map_err(Error::from)?;
1863 self.constant_from(uploaded)
1864 }
1865
1866 /// Import a trainable leaf within this borrowed session.
1867 ///
1868 /// # Examples
1869 ///
1870 /// ```rust
1871 /// use tenferro_ad::{EagerRuntime, Tensor};
1872 /// let ctx = EagerRuntime::new()?;
1873 /// let x = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |session| {
1874 /// session.variable_from(Tensor::from_vec_col_major(vec![1], vec![2.0_f64])?)
1875 /// })?;
1876 /// assert!(x.tracks_grad());
1877 /// # Ok::<(), tenferro_ad::Error>(())
1878 /// ```
1879 ///
1880 /// # Errors
1881 ///
1882 /// Returns [`Error::TensorRuntime`] for a typed backend failure when the value cannot be
1883 /// registered, or [`Error::RuntimeState`] when the runtime's value or
1884 /// gradient registry is unavailable.
1885 pub fn variable_from(&mut self, tensor: Tensor) -> Result<EagerTensor> {
1886 EagerTensor::new_leaf_in_session(Arc::clone(self.runtime), tensor, true, &mut *self.engine)
1887 }
1888
1889 /// Add eager tensors with the same broadcast and AD rules as the eager
1890 /// operation surface, reusing this borrowed execution session.
1891 ///
1892 /// # Examples
1893 ///
1894 /// ```rust
1895 /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
1896 /// let ctx = EagerRuntime::new()?;
1897 /// let x = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0])?, ctx.clone())?;
1898 /// let scalar = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![], vec![3.0_f64])?, ctx.clone())?;
1899 /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |session| session.add(&x, &scalar))?;
1900 /// assert_eq!(y.value()?.as_slice::<f64>()?, &[4.0, 5.0]);
1901 /// # Ok::<(), tenferro_ad::Error>(())
1902 /// ```
1903 ///
1904 /// # Errors
1905 ///
1906 /// Returns [`Error::ContextMismatch`] for a tensor from another runtime,
1907 /// or a typed broadcast/backend error for the operands.
1908 pub fn add(&mut self, lhs: &EagerTensor, rhs: &EagerTensor) -> Result<EagerTensor> {
1909 self.run_binary("add", lhs, rhs, StdTensorOp::Add)
1910 }
1911
1912 /// Subtract eager tensors within this borrowed session.
1913 ///
1914 /// # Examples
1915 ///
1916 /// ```rust
1917 /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
1918 /// let ctx = EagerRuntime::new()?;
1919 /// let x = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![1], vec![3.0_f64])?, ctx.clone())?;
1920 /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |session| session.sub(&x, &x))?;
1921 /// assert_eq!(y.value()?.as_slice::<f64>()?, &[0.0]);
1922 /// # Ok::<(), tenferro_ad::Error>(())
1923 /// ```
1924 ///
1925 /// # Errors
1926 ///
1927 /// Returns [`Error::ContextMismatch`] for a foreign runtime, or a typed
1928 /// broadcast/backend error for the operands.
1929 pub fn sub(&mut self, lhs: &EagerTensor, rhs: &EagerTensor) -> Result<EagerTensor> {
1930 self.run_binary("sub", lhs, rhs, StdTensorOp::Sub)
1931 }
1932
1933 /// Multiply eager tensors within this borrowed session.
1934 ///
1935 /// # Examples
1936 ///
1937 /// ```rust
1938 /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
1939 /// let ctx = EagerRuntime::new()?;
1940 /// let x = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![1], vec![3.0_f64])?, ctx.clone())?;
1941 /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |session| session.mul(&x, &x))?;
1942 /// assert_eq!(y.value()?.as_slice::<f64>()?, &[9.0]);
1943 /// # Ok::<(), tenferro_ad::Error>(())
1944 /// ```
1945 ///
1946 /// # Errors
1947 ///
1948 /// Returns [`Error::ContextMismatch`] for a foreign runtime, or a typed
1949 /// broadcast/backend error for the operands.
1950 pub fn mul(&mut self, lhs: &EagerTensor, rhs: &EagerTensor) -> Result<EagerTensor> {
1951 self.run_binary("mul", lhs, rhs, StdTensorOp::Mul)
1952 }
1953
1954 /// Divide eager tensors elementwise with broadcast rules.
1955 ///
1956 /// # Examples
1957 /// ```rust
1958 /// use tenferro_ad::{EagerRuntime, Tensor};
1959 /// let ctx = EagerRuntime::new()?;
1960 /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
1961 /// let x = s.constant_from(Tensor::from_vec_col_major(vec![1], vec![6.0_f64])?)?;
1962 /// let divisor = s.constant_from(Tensor::from_vec_col_major(vec![], vec![2.0_f64])?)?;
1963 /// s.div(&x, &divisor)
1964 /// })?;
1965 /// assert_eq!(y.value()?.as_slice::<f64>()?, &[3.0]);
1966 /// # Ok::<(), tenferro_ad::Error>(())
1967 /// ```
1968 /// # Errors
1969 /// Returns [`Error::ContextMismatch`] when an input belongs to another eager
1970 /// runtime, a validation error with
1971 /// `ValidationError::ShapeMismatch` when the operands cannot broadcast, or
1972 /// [`Error::TensorRuntime`] for a typed backend failure (including integer division by
1973 /// zero).
1974 pub fn div(&mut self, lhs: &EagerTensor, rhs: &EagerTensor) -> Result<EagerTensor> {
1975 self.run_binary("div", lhs, rhs, StdTensorOp::Div)
1976 }
1977
1978 /// Compute the elementwise remainder with broadcast rules.
1979 ///
1980 /// # Examples
1981 /// ```rust
1982 /// use tenferro_ad::{EagerRuntime, Tensor};
1983 /// let ctx = EagerRuntime::new()?;
1984 /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
1985 /// let x = s.constant_from(Tensor::from_vec_col_major(vec![1], vec![5.0_f64])?)?;
1986 /// let divisor = s.constant_from(Tensor::from_vec_col_major(vec![], vec![2.0_f64])?)?;
1987 /// s.rem(&x, &divisor)
1988 /// })?;
1989 /// assert_eq!(y.value()?.as_slice::<f64>()?, &[1.0]);
1990 /// # Ok::<(), tenferro_ad::Error>(())
1991 /// ```
1992 /// # Errors
1993 /// Returns [`Error::ContextMismatch`] when an input belongs to another eager
1994 /// runtime, a validation error with
1995 /// `ValidationError::ShapeMismatch` when the operands cannot broadcast, or
1996 /// [`Error::TensorRuntime`] for a typed backend failure (including an integer remainder by
1997 /// zero).
1998 pub fn rem(&mut self, lhs: &EagerTensor, rhs: &EagerTensor) -> Result<EagerTensor> {
1999 self.run_binary("rem", lhs, rhs, StdTensorOp::Rem)
2000 }
2001
2002 /// Raise eager tensor elements to broadcast exponents.
2003 ///
2004 /// # Examples
2005 /// ```rust
2006 /// use tenferro_ad::{EagerRuntime, Tensor};
2007 /// let ctx = EagerRuntime::new()?;
2008 /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
2009 /// let x = s.constant_from(Tensor::from_vec_col_major(vec![1], vec![2.0_f64])?)?;
2010 /// let exponent = s.constant_from(Tensor::from_vec_col_major(vec![], vec![3.0_f64])?)?;
2011 /// s.pow(&x, &exponent)
2012 /// })?;
2013 /// assert_eq!(y.value()?.as_slice::<f64>()?, &[8.0]);
2014 /// # Ok::<(), tenferro_ad::Error>(())
2015 /// ```
2016 /// # Errors
2017 /// Returns [`Error::ContextMismatch`] when an input belongs to another eager
2018 /// runtime, a validation error with
2019 /// `ValidationError::ShapeMismatch` when the operands cannot broadcast, or
2020 /// [`Error::TensorRuntime`] for a typed backend failure (including a negative integer
2021 /// exponent).
2022 pub fn pow(&mut self, lhs: &EagerTensor, rhs: &EagerTensor) -> Result<EagerTensor> {
2023 self.run_binary("pow", lhs, rhs, StdTensorOp::Pow)
2024 }
2025
2026 /// Compute the elementwise maximum under broadcast rules.
2027 ///
2028 /// # Examples
2029 /// ```rust
2030 /// use tenferro_ad::{EagerRuntime, Tensor};
2031 /// let ctx = EagerRuntime::new()?;
2032 /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
2033 /// let x = s.constant_from(Tensor::from_vec_col_major(vec![1], vec![2.0_f64])?)?;
2034 /// let bound = s.constant_from(Tensor::from_vec_col_major(vec![], vec![3.0_f64])?)?;
2035 /// s.maximum(&x, &bound)
2036 /// })?;
2037 /// assert_eq!(y.value()?.as_slice::<f64>()?, &[3.0]);
2038 /// # Ok::<(), tenferro_ad::Error>(())
2039 /// ```
2040 /// # Errors
2041 /// Returns a typed foreign-runtime, broadcast, unsupported-dtype, or backend error.
2042 pub fn maximum(&mut self, lhs: &EagerTensor, rhs: &EagerTensor) -> Result<EagerTensor> {
2043 self.run_binary("maximum", lhs, rhs, StdTensorOp::Maximum)
2044 }
2045
2046 /// Compute the elementwise minimum under broadcast rules.
2047 ///
2048 /// # Examples
2049 /// ```rust
2050 /// use tenferro_ad::{EagerRuntime, Tensor};
2051 /// let ctx = EagerRuntime::new()?;
2052 /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
2053 /// let x = s.constant_from(Tensor::from_vec_col_major(vec![1], vec![2.0_f64])?)?;
2054 /// let bound = s.constant_from(Tensor::from_vec_col_major(vec![], vec![3.0_f64])?)?;
2055 /// s.minimum(&x, &bound)
2056 /// })?;
2057 /// assert_eq!(y.value()?.as_slice::<f64>()?, &[2.0]);
2058 /// # Ok::<(), tenferro_ad::Error>(())
2059 /// ```
2060 /// # Errors
2061 /// Returns a typed foreign-runtime, broadcast, unsupported-dtype, or backend error.
2062 pub fn minimum(&mut self, lhs: &EagerTensor, rhs: &EagerTensor) -> Result<EagerTensor> {
2063 self.run_binary("minimum", lhs, rhs, StdTensorOp::Minimum)
2064 }
2065
2066 /// Compare eager tensors elementwise under broadcast rules.
2067 ///
2068 /// # Examples
2069 /// ```rust
2070 /// use tenferro_ad::{CompareDir, EagerRuntime, Tensor};
2071 /// let ctx = EagerRuntime::new()?;
2072 /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
2073 /// let x = s.constant_from(Tensor::from_vec_col_major(vec![1], vec![2.0_f64])?)?;
2074 /// let bound = s.constant_from(Tensor::from_vec_col_major(vec![], vec![1.0_f64])?)?;
2075 /// s.compare(&x, &bound, CompareDir::Gt)
2076 /// })?;
2077 /// assert_eq!(y.value()?.as_slice::<bool>()?, &[true]);
2078 /// # Ok::<(), tenferro_ad::Error>(())
2079 /// ```
2080 /// # Errors
2081 /// Returns a typed foreign-runtime, broadcast, unsupported-dtype, or backend error.
2082 pub fn compare(
2083 &mut self,
2084 lhs: &EagerTensor,
2085 rhs: &EagerTensor,
2086 dir: CompareDir,
2087 ) -> Result<EagerTensor> {
2088 self.run_binary("compare", lhs, rhs, StdTensorOp::Compare(dir))
2089 }
2090
2091 /// Select eager values elementwise using a broadcast boolean condition.
2092 ///
2093 /// # Examples
2094 /// ```rust
2095 /// use tenferro_ad::{EagerRuntime, Tensor};
2096 /// let ctx = EagerRuntime::new()?;
2097 /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
2098 /// let condition = s.constant_from(Tensor::from_vec_col_major(vec![2], vec![true, false])?)?;
2099 /// let yes = s.constant_from(Tensor::from_vec_col_major(vec![], vec![10.0_f64])?)?;
2100 /// let no = s.constant_from(Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0])?)?;
2101 /// s.where_select(&condition, &yes, &no)
2102 /// })?;
2103 /// assert_eq!(y.value()?.as_slice::<f64>()?, &[10.0, 2.0]);
2104 /// # Ok::<(), tenferro_ad::Error>(())
2105 /// ```
2106 /// # Errors
2107 /// Returns a typed foreign-runtime, broadcast, unsupported-dtype, or backend error.
2108 pub fn where_select(
2109 &mut self,
2110 condition: &EagerTensor,
2111 on_true: &EagerTensor,
2112 on_false: &EagerTensor,
2113 ) -> Result<EagerTensor> {
2114 self.run_ternary(
2115 "where_select",
2116 condition,
2117 on_true,
2118 on_false,
2119 StdTensorOp::Select,
2120 )
2121 }
2122
2123 /// Alias for [`Self::where_select`] with the same borrowed-session semantics.
2124 ///
2125 /// # Examples
2126 /// ```rust
2127 /// use tenferro_ad::{EagerRuntime, Tensor};
2128 /// let ctx = EagerRuntime::new()?;
2129 /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
2130 /// let predicate = s.constant_from(Tensor::from_vec_col_major(vec![], vec![true])?)?;
2131 /// let yes = s.constant_from(Tensor::from_vec_col_major(vec![], vec![3.0_f64])?)?;
2132 /// let no = s.constant_from(Tensor::from_vec_col_major(vec![], vec![4.0_f64])?)?;
2133 /// s.select(&predicate, &yes, &no)
2134 /// })?;
2135 /// assert_eq!(y.value()?.as_slice::<f64>()?, &[3.0]);
2136 /// # Ok::<(), tenferro_ad::Error>(())
2137 /// ```
2138 /// # Errors
2139 /// Returns a typed foreign-runtime, broadcast, unsupported-dtype, or backend error.
2140 pub fn select(
2141 &mut self,
2142 condition: &EagerTensor,
2143 on_true: &EagerTensor,
2144 on_false: &EagerTensor,
2145 ) -> Result<EagerTensor> {
2146 self.where_select(condition, on_true, on_false)
2147 }
2148
2149 /// Clamp eager values elementwise between broadcast lower and upper bounds.
2150 ///
2151 /// # Examples
2152 /// ```rust
2153 /// use tenferro_ad::{EagerRuntime, Tensor};
2154 /// let ctx = EagerRuntime::new()?;
2155 /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
2156 /// let x = s.constant_from(Tensor::from_vec_col_major(vec![2], vec![-2.0_f64, 5.0])?)?;
2157 /// let lo = s.constant_from(Tensor::from_vec_col_major(vec![], vec![-1.0_f64])?)?;
2158 /// let hi = s.constant_from(Tensor::from_vec_col_major(vec![], vec![4.0_f64])?)?;
2159 /// s.clamp(&x, &lo, &hi)
2160 /// })?;
2161 /// assert_eq!(y.value()?.as_slice::<f64>()?, &[-1.0, 4.0]);
2162 /// # Ok::<(), tenferro_ad::Error>(())
2163 /// ```
2164 /// # Errors
2165 /// Returns a typed foreign-runtime, broadcast, unsupported-dtype, or backend error.
2166 pub fn clamp(
2167 &mut self,
2168 input: &EagerTensor,
2169 lower: &EagerTensor,
2170 upper: &EagerTensor,
2171 ) -> Result<EagerTensor> {
2172 self.run_ternary("clamp", input, lower, upper, StdTensorOp::Clamp)
2173 }
2174
2175 /// Contract eager tensors according to a dot-general dimension mapping.
2176 ///
2177 /// The output layout is `[lhs free..., rhs free..., batch...]`: batch axes
2178 /// come last (see [`DotGeneralConfig`](tenferro_runtime::DotGeneralConfig)).
2179 ///
2180 /// # Examples
2181 ///
2182 /// ```rust
2183 /// use tenferro_ad::{DotGeneralConfig, EagerRuntime, Tensor};
2184 /// let ctx = EagerRuntime::new()?;
2185 /// let result = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |session| {
2186 /// let lhs = session.variable_from(Tensor::from_vec_col_major(vec![1, 2], vec![2.0_f64, 3.0])?)?;
2187 /// let rhs = session.constant_from(Tensor::from_vec_col_major(vec![2, 1], vec![4.0_f64, 5.0])?)?;
2188 /// session.dot_general(&lhs, &rhs, DotGeneralConfig {
2189 /// lhs_contracting_dims: [1].as_slice().into(),
2190 /// rhs_contracting_dims: [0].as_slice().into(),
2191 /// lhs_batch_dims: [].as_slice().into(),
2192 /// rhs_batch_dims: [].as_slice().into(),
2193 /// })
2194 /// })?;
2195 /// assert_eq!(result.value()?.as_slice::<f64>()?, &[23.0]);
2196 /// # Ok::<(), tenferro_ad::Error>(())
2197 /// ```
2198 ///
2199 /// # Errors
2200 ///
2201 /// Returns [`Error::ContextMismatch`] for a foreign eager runtime,
2202 /// a typed validation error for incompatible contraction dimensions,
2203 /// or the backend's typed execution error.
2204 pub fn dot_general(
2205 &mut self,
2206 lhs: &EagerTensor,
2207 rhs: &EagerTensor,
2208 config: DotGeneralConfig,
2209 ) -> Result<EagerTensor> {
2210 self.ensure_runtime(lhs)?;
2211 self.ensure_runtime(rhs)?;
2212 config
2213 .validate_dims_with_ranks(lhs.shape().len(), rhs.shape().len())
2214 .map_err(Error::TensorRuntime)?;
2215 EagerTensor::nary_op_in_session(
2216 &[lhs, rhs],
2217 StdTensorOp::DotGeneral { config },
2218 &mut *self.engine,
2219 )
2220 }
2221
2222 /// Scale an eager tensor by a real scalar in this borrowed session.
2223 /// Integer factors are rounded; finite zero maps to `false` for boolean inputs.
2224 ///
2225 /// # Examples
2226 /// ```rust
2227 /// use tenferro_ad::{EagerRuntime, Tensor};
2228 /// let ctx = EagerRuntime::new()?;
2229 /// let scaled = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
2230 /// let x = s.constant_from(Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0])?)?;
2231 /// s.scale_real(&x, 2.0)
2232 /// })?;
2233 /// assert_eq!(scaled.value()?.as_slice::<f64>()?, &[2.0, 4.0]);
2234 /// # Ok::<(), tenferro_ad::Error>(())
2235 /// ```
2236 /// # Errors
2237 /// Returns a typed foreign-runtime, invalid-factor/dtype, or backend error.
2238 pub fn scale_real(&mut self, input: &EagerTensor, factor: f64) -> Result<EagerTensor> {
2239 self.ensure_runtime(input)?;
2240 let scalar = tenferro_runtime::scale::real_scale_scalar(input.dtype(), factor)?;
2241 let scalar = self.constant_from_host(scalar)?;
2242 self.mul(input, &scalar)
2243 }
2244
2245 /// Scale a complex eager tensor by a complex scalar in this borrowed session.
2246 ///
2247 /// # Examples
2248 /// ```rust
2249 /// use num_complex::Complex64;
2250 /// use tenferro_ad::{EagerRuntime, Tensor};
2251 /// let ctx = EagerRuntime::new()?;
2252 /// let scaled = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
2253 /// let x = s.constant_from(Tensor::from_vec_col_major(vec![1], vec![Complex64::new(1.0, 2.0)])?)?;
2254 /// s.scale_complex(&x, Complex64::new(0.0, 1.0))
2255 /// })?;
2256 /// assert_eq!(scaled.value()?.as_slice::<Complex64>()?, &[Complex64::new(-2.0, 1.0)]);
2257 /// # Ok::<(), tenferro_ad::Error>(())
2258 /// ```
2259 /// # Errors
2260 /// Returns [`Error::ContextMismatch`] when an input belongs to another eager
2261 /// runtime, [`Error::TensorRuntime`] containing
2262 /// `ValidationError::InvalidArgument` when the input dtype is not complex,
2263 /// or [`Error::TensorRuntime`] for a typed backend failure.
2264 pub fn scale_complex(&mut self, input: &EagerTensor, factor: Complex64) -> Result<EagerTensor> {
2265 self.ensure_runtime(input)?;
2266 let scalar = tenferro_runtime::scale::complex_scale_scalar(input.dtype(), factor)?;
2267 let scalar = self.constant_from_host(scalar)?;
2268 self.mul(input, &scalar)
2269 }
2270
2271 /// Multiply two rank-2 eager tensors in this borrowed session.
2272 ///
2273 /// # Examples
2274 /// ```rust
2275 /// use tenferro_ad::{EagerRuntime, Tensor};
2276 /// let ctx = EagerRuntime::new()?;
2277 /// let result = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
2278 /// let a = s.constant_from(Tensor::from_vec_col_major(vec![1, 1], vec![2.0_f64])?)?;
2279 /// let b = s.constant_from(Tensor::from_vec_col_major(vec![1, 1], vec![3.0_f64])?)?;
2280 /// s.matmul(&a, &b)
2281 /// })?;
2282 /// assert_eq!(result.value()?.as_slice::<f64>()?, &[6.0]);
2283 /// # Ok::<(), tenferro_ad::Error>(())
2284 /// ```
2285 /// # Errors
2286 /// Returns [`Error::ContextMismatch`] when an input belongs to another eager
2287 /// runtime, a validation error with
2288 /// `ValidationError::RankMismatch` or `ValidationError::ShapeMismatch` when
2289 /// the operands are not rank-2 with matching inner dimensions, a dtype
2290 /// mismatch between the operands, or [`Error::TensorRuntime`] for a typed backend failure.
2291 pub fn matmul(&mut self, lhs: &EagerTensor, rhs: &EagerTensor) -> Result<EagerTensor> {
2292 self.ensure_runtime(lhs)?;
2293 self.ensure_runtime(rhs)?;
2294 let lhs_shape = lhs.shape();
2295 let rhs_shape = rhs.shape();
2296 if lhs_shape.len() != 2 {
2297 return Err(tenferro_tensor::Error::rank_mismatch("matmul", 2, lhs_shape.len()).into());
2298 }
2299 if rhs_shape.len() != 2 {
2300 return Err(tenferro_tensor::Error::rank_mismatch("matmul", 2, rhs_shape.len()).into());
2301 }
2302 if lhs_shape[1] != rhs_shape[0] {
2303 return Err(
2304 tenferro_tensor::Error::shape_mismatch("matmul", lhs_shape, rhs_shape).into(),
2305 );
2306 }
2307 self.dot_general(
2308 lhs,
2309 rhs,
2310 DotGeneralConfig {
2311 lhs_contracting_dims: [1].as_slice().into(),
2312 rhs_contracting_dims: [0].as_slice().into(),
2313 lhs_batch_dims: [].as_slice().into(),
2314 rhs_batch_dims: [].as_slice().into(),
2315 },
2316 )
2317 }
2318
2319 /// Contract eagerly with optional conjugation of either operand.
2320 /// Untracked operands use the backend's conjugating contraction directly;
2321 /// tracked operands record explicit conjugations for reverse-mode AD.
2322 ///
2323 /// # Examples
2324 ///
2325 /// ```rust
2326 /// use tenferro_ad::{DotGeneralConfig, EagerRuntime, Tensor};
2327 /// let ctx = EagerRuntime::new()?;
2328 /// let result = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |session| {
2329 /// let lhs = session.constant_from(Tensor::from_vec_col_major(vec![1, 1], vec![2.0_f64])?)?;
2330 /// let rhs = session.constant_from(Tensor::from_vec_col_major(vec![1, 1], vec![3.0_f64])?)?;
2331 /// session.dot_general_with_conj(&lhs, &rhs, DotGeneralConfig {
2332 /// lhs_contracting_dims: [1].as_slice().into(),
2333 /// rhs_contracting_dims: [0].as_slice().into(),
2334 /// lhs_batch_dims: [].as_slice().into(),
2335 /// rhs_batch_dims: [].as_slice().into(),
2336 /// }, true, false)
2337 /// })?;
2338 /// assert_eq!(result.value()?.as_slice::<f64>()?, &[6.0]);
2339 /// # Ok::<(), tenferro_ad::Error>(())
2340 /// ```
2341 ///
2342 /// # Errors
2343 ///
2344 /// Returns [`Error::ContextMismatch`] for a foreign eager runtime,
2345 /// a typed validation error for invalid dimensions, or a backend error.
2346 pub fn dot_general_with_conj(
2347 &mut self,
2348 lhs: &EagerTensor,
2349 rhs: &EagerTensor,
2350 config: DotGeneralConfig,
2351 lhs_conj: bool,
2352 rhs_conj: bool,
2353 ) -> Result<EagerTensor> {
2354 self.ensure_runtime(lhs)?;
2355 self.ensure_runtime(rhs)?;
2356 config
2357 .validate_dims_with_ranks(lhs.shape().len(), rhs.shape().len())
2358 .map_err(Error::TensorRuntime)?;
2359 if !lhs.requires_grad && !rhs.requires_grad {
2360 let output = crate::eager_exec::exec_dot_general_with_conj_on_tensor_reads_in_session(
2361 lhs.tensor_read(),
2362 rhs.tensor_read(),
2363 &config,
2364 lhs_conj,
2365 rhs_conj,
2366 &mut *self.engine,
2367 )?;
2368 return EagerTensor::new_untracked_result(Arc::clone(self.runtime), output);
2369 }
2370 let lhs = if lhs_conj {
2371 self.conj(lhs)?
2372 } else {
2373 lhs.clone()
2374 };
2375 let rhs = if rhs_conj {
2376 self.conj(rhs)?
2377 } else {
2378 rhs.clone()
2379 };
2380 self.dot_general(&lhs, &rhs, config)
2381 }
2382
2383 fn run_binary(
2384 &mut self,
2385 name: &'static str,
2386 lhs: &EagerTensor,
2387 rhs: &EagerTensor,
2388 op: StdTensorOp,
2389 ) -> Result<EagerTensor> {
2390 self.ensure_runtime(lhs)?;
2391 self.ensure_runtime(rhs)?;
2392 let (lhs, rhs) = crate::eager_ops::broadcast_binary_in_session(name, lhs, rhs, self)?;
2393 EagerTensor::nary_op_in_session(&[lhs.as_ref(), rhs.as_ref()], op, &mut *self.engine)
2394 }
2395
2396 fn run_ternary(
2397 &mut self,
2398 name: &'static str,
2399 first: &EagerTensor,
2400 second: &EagerTensor,
2401 third: &EagerTensor,
2402 op: StdTensorOp,
2403 ) -> Result<EagerTensor> {
2404 self.ensure_runtime(first)?;
2405 self.ensure_runtime(second)?;
2406 self.ensure_runtime(third)?;
2407 let (first, second, third) =
2408 crate::eager_ops::broadcast_ternary_in_session(name, first, second, third, self)?;
2409 EagerTensor::nary_op_in_session(&[&first, &second, &third], op, &mut *self.engine)
2410 }
2411
2412 /// Apply one standard tensor op in this borrowed session and record it
2413 /// for AD when needed.
2414 ///
2415 /// Extension crates use this when an extension-level eager operation
2416 /// expands into ordinary `StdTensorOp` nodes instead of a custom extension
2417 /// primitive: all of them run in this one backend session.
2418 ///
2419 /// # Examples
2420 ///
2421 /// ```rust
2422 /// use tenferro_ad::{EagerRuntime, Tensor};
2423 /// use tenferro_ops::std_tensor_op::StdTensorOp;
2424 ///
2425 /// let ctx = EagerRuntime::new()?;
2426 /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
2427 /// let x = s.variable_from(Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0])?)?;
2428 /// let negated = s.apply_standard_op(StdTensorOp::Neg, &[&x])?;
2429 /// s.apply_standard_op(StdTensorOp::Mul, &[&negated, &x])
2430 /// })?;
2431 /// assert_eq!(y.value()?.as_slice::<f64>()?, &[-1.0, -4.0]);
2432 /// assert!(y.tracks_grad());
2433 /// # Ok::<(), tenferro_ad::Error>(())
2434 /// ```
2435 ///
2436 /// # Errors
2437 ///
2438 /// Returns [`Error::TensorRuntime`] containing
2439 /// [`tenferro_tensor::ValidationError::InvalidArgument`] for an extension
2440 /// op, [`Error::ContextMismatch`] for a tensor from another runtime, a
2441 /// typed input-count error, or the backend's typed execution error.
2442 pub fn apply_standard_op(
2443 &mut self,
2444 op: StdTensorOp,
2445 inputs: &[&EagerTensor],
2446 ) -> Result<EagerTensor> {
2447 if matches!(op, StdTensorOp::Extension(_)) {
2448 return Err(Error::invalid_argument(
2449 "EagerSession::apply_standard_op",
2450 ErrorPhase::Execution,
2451 "op",
2452 "Extension ops must be passed to apply_eager",
2453 ));
2454 }
2455 for input in inputs {
2456 self.ensure_runtime(input)?;
2457 }
2458 EagerTensor::nary_op_in_session(inputs, op, &mut *self.engine)
2459 }
2460
2461 /// Borrow the backend session this eager session runs on.
2462 ///
2463 /// Extension crates use it to run their backend kernels on untracked
2464 /// values inside the same execution region instead of entering a second
2465 /// session, which would be rejected as reentry. It grants the same access
2466 /// as passing the engine directly.
2467 ///
2468 /// # Examples
2469 ///
2470 /// ```rust
2471 /// use tenferro_ad::{EagerRuntime, Tensor};
2472 /// use tenferro_tensor::TensorRead;
2473 ///
2474 /// let ctx = EagerRuntime::new()?;
2475 /// let x = Tensor::from_vec_col_major(vec![2], vec![1.0_f64, -2.0])?;
2476 /// let copy = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
2477 /// s.backend_session()
2478 /// .to_contiguous_read(TensorRead::from_tensor(&x))
2479 /// .map_err(tenferro_ad::Error::from)
2480 /// })?;
2481 /// assert_eq!(copy.as_slice::<f64>()?, &[1.0, -2.0]);
2482 /// # Ok::<(), Box<dyn std::error::Error>>(())
2483 /// ```
2484 pub fn backend_session(&mut self) -> &mut dyn Engine {
2485 &mut *self.engine
2486 }
2487
2488 /// Borrow the engine this eager session runs on.
2489 ///
2490 /// # Examples
2491 ///
2492 /// ```rust
2493 /// use tenferro_ad::EagerRuntime;
2494 /// use tenferro_cpu::CpuEngine;
2495 ///
2496 /// let ctx = EagerRuntime::new()?;
2497 /// let mut engine = CpuEngine::serial();
2498 /// let id = ctx.with_eager_session(&mut engine, |s| {
2499 /// Ok::<_, tenferro_ad::Error>(s.engine().target().engine_id().clone())
2500 /// })?;
2501 /// assert_eq!(id, tenferro_cpu::runtime_engine_id()?);
2502 /// # Ok::<(), Box<dyn std::error::Error>>(())
2503 /// ```
2504 pub fn engine(&mut self) -> &mut dyn ExecutionEngine {
2505 &mut *self.engine
2506 }
2507
2508 fn run_unary(&mut self, input: &EagerTensor, op: StdTensorOp) -> Result<EagerTensor> {
2509 self.ensure_runtime(input)?;
2510 EagerTensor::nary_op_in_session(&[input], op, &mut *self.engine)
2511 }
2512
2513 pub(crate) fn record_outputs(
2514 &mut self,
2515 op: &StdTensorOp,
2516 outputs: &[&Tensor],
2517 inputs: &[&EagerTensor],
2518 ) -> Result<RecordedEagerOutputs> {
2519 record_eager_outputs_in_session(op, outputs, inputs, &mut *self.engine)
2520 }
2521
2522 pub(crate) fn ensure_runtime(&self, input: &EagerTensor) -> Result<()> {
2523 if !Arc::ptr_eq(self.runtime, &input.ctx) {
2524 return Err(Error::ContextMismatch {
2525 lhs: self.runtime.id(),
2526 rhs: input.ctx_id(),
2527 });
2528 }
2529 Ok(())
2530 }
2531
2532 pub(crate) fn runtime(&self) -> &Arc<EagerRuntime> {
2533 self.runtime
2534 }
2535
2536 /// Run `f` on this runtime's extension cache store from inside the
2537 /// session.
2538 ///
2539 /// Operation families use this for their own prepared-plan caches without
2540 /// reopening the runtime: the eager owner is already locked, and the cache
2541 /// lock is taken second, as in every extension execution region.
2542 ///
2543 /// # Examples
2544 ///
2545 /// ```rust
2546 /// use tenferro_ad::EagerRuntime;
2547 ///
2548 /// let ctx = EagerRuntime::new()?;
2549 /// let entries = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |session| {
2550 /// session.with_extension_caches(|caches| caches.len())
2551 /// })?;
2552 /// assert_eq!(entries, 0);
2553 /// # Ok::<(), tenferro_ad::Error>(())
2554 /// ```
2555 ///
2556 /// # Errors
2557 ///
2558 /// Returns a runtime-state error when the extension cache lock is
2559 /// poisoned.
2560 pub fn with_extension_caches<R>(
2561 &mut self,
2562 f: impl FnOnce(&mut tenferro_runtime::ExtensionCacheStore) -> R,
2563 ) -> Result<R> {
2564 let result = with_lent_extension_caches(self.engine, |_, caches| f(caches));
2565 Ok(result)
2566 }
2567}
2568
2569impl fmt::Debug for EagerRuntime {
2570 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
2571 let mut debug = f.debug_struct("EagerRuntime");
2572 debug.field("id", &self.id);
2573 debug.field("runtime_id", &self.runtime.id());
2574 debug.field("runtime_epoch", &self.runtime.epoch().ok());
2575 match self.extension_install_lock.try_lock() {
2576 Ok(_) => {
2577 debug.field("extension_install_lock", &"<unlocked>");
2578 }
2579 Err(_) => {
2580 debug.field("extension_install_lock", &"<locked>");
2581 }
2582 }
2583 debug.field("semantic_extension_rules", &self.semantic_extension_rules);
2584 match self.grad_slots.try_lock() {
2585 Ok(slots) => {
2586 debug.field("grad_slots_len", &slots.len());
2587 }
2588 Err(_) => {
2589 debug.field("grad_slots_len", &"<locked>");
2590 }
2591 }
2592 match self.value_records.try_lock() {
2593 Ok(records) => {
2594 debug.field("value_records_len", &records.len());
2595 }
2596 Err(_) => {
2597 debug.field("value_records_len", &"<locked>");
2598 }
2599 }
2600 match self.ad_transform_cache.stats() {
2601 Ok(stats) => {
2602 debug.field("ad_transform_cache_stats", &stats);
2603 }
2604 Err(err) => {
2605 debug.field("ad_transform_cache_stats", &format_args!("{err}"));
2606 }
2607 }
2608 match self.prepared_derivative_cache.try_lock() {
2609 Ok(cache) => {
2610 debug.field("prepared_derivative_cache_stats", &cache.stats());
2611 }
2612 Err(_) => {
2613 debug.field("prepared_derivative_cache_stats", &"<locked>");
2614 }
2615 }
2616 debug.finish_non_exhaustive()
2617 }
2618}
2619
2620impl EagerRuntime {
2621 fn lock_extension_install(&self) -> Result<MutexGuard<'_, ()>> {
2622 self.extension_install_lock.lock().map_err(|_| {
2623 Error::runtime_state(
2624 "eager_extension_install",
2625 ErrorPhase::Execution,
2626 "lock poisoned",
2627 )
2628 })
2629 }
2630
2631 fn lock_prepared_derivative_cache(&self) -> Result<MutexGuard<'_, PreparedDerivativeCache>> {
2632 self.prepared_derivative_cache.lock().map_err(|_| {
2633 Error::runtime_state(
2634 "prepared_derivative_cache",
2635 ErrorPhase::Execution,
2636 "lock poisoned",
2637 )
2638 })
2639 }
2640
2641 fn lock_grad_slots(
2642 &self,
2643 ) -> Result<MutexGuard<'_, HashMap<ValueKey<StdTensorOp>, WeakGradSlot>>> {
2644 self.grad_slots.lock().map_err(|_| {
2645 Error::runtime_state(
2646 "eager_gradient_slots",
2647 ErrorPhase::Execution,
2648 "lock poisoned",
2649 )
2650 })
2651 }
2652
2653 fn lock_value_records(
2654 &self,
2655 ) -> Result<MutexGuard<'_, HashMap<ValueKey<StdTensorOp>, Weak<EagerTensorRecord>>>> {
2656 self.value_records.lock().map_err(|_| {
2657 Error::runtime_state(
2658 "eager_value_registry",
2659 ErrorPhase::Execution,
2660 "lock poisoned",
2661 )
2662 })
2663 }
2664
2665 fn with_rules_and_cache(
2666 semantic_extension_rules: SemanticExtensionRuleSet,
2667 ad_transform_cache: Arc<AdTransformCache>,
2668 ) -> Result<Self> {
2669 let runtime = Runtime::builder()
2670 .build()
2671 .map_err(|source| runtime_config_error("EagerRuntime::new", source))?;
2672 Ok(Self {
2673 id: ContextId::fresh(),
2674 runtime,
2675 extension_install_lock: Mutex::new(()),
2676 semantic_extension_rules,
2677 grad_slots: Mutex::new(HashMap::new()),
2678 value_records: Mutex::new(HashMap::new()),
2679 ad_transform_cache,
2680 prepared_derivative_cache: Mutex::new(PreparedDerivativeCache::default()),
2681 })
2682 }
2683
2684 /// Create a shared eager execution context.
2685 ///
2686 /// The context stores no engine: every operation runs on the engine the
2687 /// caller passes in.
2688 ///
2689 /// # Examples
2690 ///
2691 /// ```
2692 /// use tenferro_ad::EagerRuntime;
2693 ///
2694 /// let ctx = EagerRuntime::new()?;
2695 /// assert_eq!(std::sync::Arc::strong_count(&ctx), 1);
2696 /// # Ok::<(), tenferro_ad::Error>(())
2697 /// ```
2698 ///
2699 /// # Errors
2700 ///
2701 /// Returns [`Error::RuntimeStateSource`] when the prepared-program runtime
2702 /// cannot be configured, preserving the underlying [`RuntimeConfigError`].
2703 pub fn new() -> Result<Arc<Self>> {
2704 Ok(Arc::new(Self::with_rules_and_cache(
2705 SemanticExtensionRuleSet::default(),
2706 Arc::new(AdTransformCache::new()),
2707 )?))
2708 }
2709
2710 /// Create a shared eager execution context that uses the semantic
2711 /// extension rules and AD transform cache of `ad`.
2712 ///
2713 /// # Examples
2714 ///
2715 /// ```
2716 /// use tenferro_ad::{AdContext, EagerRuntime};
2717 ///
2718 /// let ad = AdContext::builder().build().unwrap();
2719 /// let ctx = EagerRuntime::with_ad_context(&ad)?;
2720 /// assert_eq!(std::sync::Arc::strong_count(&ctx), 1);
2721 /// # Ok::<(), tenferro_ad::Error>(())
2722 /// ```
2723 ///
2724 /// # Errors
2725 ///
2726 /// Returns [`Error::RuntimeStateSource`] when the prepared-program runtime
2727 /// cannot be configured.
2728 pub fn with_ad_context(ad: &AdContext) -> Result<Arc<Self>> {
2729 Ok(Arc::new(Self::with_rules_and_cache(
2730 ad.semantic_extension_rules().clone(),
2731 ad.ad_transform_cache(),
2732 )?))
2733 }
2734
2735 /// Register the kind of `engine` with this runtime if it is not registered
2736 /// yet, and return its target.
2737 pub(crate) fn ensure_engine(&self, engine: &mut dyn ExecutionEngine) -> Result<EngineTarget> {
2738 let target = engine.target();
2739 let snapshot = self
2740 .runtime
2741 .snapshot()
2742 .map_err(|source| runtime_state_source("EagerRuntime::ensure_engine", source))?;
2743 if let Some(registered) = snapshot.engine(target.engine_id()) {
2744 if registered.context_identity() != target.context_identity() {
2745 return Err(Error::unsupported(
2746 "EagerRuntime::ensure_engine",
2747 ErrorPhase::Execution,
2748 format!(
2749 "engine id {:?} is registered for a different engine type",
2750 target.engine_id()
2751 ),
2752 ));
2753 }
2754 return Ok(target);
2755 }
2756 let registration = engine_kind_registration(engine)
2757 .ok_or_else(|| unregisterable_engine_error(target.engine_id()))?;
2758 let registration = registration
2759 .map_err(|source| runtime_config_error("EagerRuntime::ensure_engine", source))?;
2760 let _install_guard = self.lock_extension_install()?;
2761 let snapshot = self
2762 .runtime
2763 .snapshot()
2764 .map_err(|source| runtime_state_source("EagerRuntime::ensure_engine", source))?;
2765 if snapshot.engine(target.engine_id()).is_none() {
2766 self.runtime
2767 .reconfigure(|edit| {
2768 edit.register_engine(registration)?;
2769 Ok(())
2770 })
2771 .map_err(|source| runtime_state_source("EagerRuntime::ensure_engine", source))?;
2772 }
2773 Ok(target)
2774 }
2775
2776 /// Return an opaque identifier for this context.
2777 ///
2778 /// # Examples
2779 ///
2780 /// ```
2781 /// use tenferro_ad::{EagerRuntime};
2782 ///
2783 /// let ctx = EagerRuntime::new()?;
2784 /// assert_ne!(ctx.id(), EagerRuntime::new()?.id());
2785 /// # Ok::<(), tenferro_ad::Error>(())
2786 /// ```
2787 pub fn id(&self) -> ContextId {
2788 self.id
2789 }
2790
2791 /// Disable eager operation recording on the current thread until the guard is dropped.
2792 ///
2793 /// This is useful for optimizer updates, metric calculations, and other
2794 /// eager computations that should not become part of the AD tape.
2795 ///
2796 /// # Examples
2797 ///
2798 /// ```
2799 /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
2800 ///
2801 /// let ctx = EagerRuntime::new()?;
2802 /// let x = EagerTensor::requires_grad_in(
2803 /// Tensor::from_vec_col_major(vec![1], vec![3.0_f64]).unwrap(),
2804 /// ctx.clone(),
2805 /// )?;
2806 /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
2807 /// let _guard = ctx.no_grad();
2808 /// s.mul(&x, &x)
2809 /// })?;
2810 /// assert!(!y.tracks_grad());
2811 /// # Ok::<(), tenferro_ad::Error>(())
2812 /// ```
2813 pub fn no_grad(&self) -> EagerNoGradGuard {
2814 EAGER_NO_GRAD_DEPTH.with(|depth| {
2815 depth.set(depth.get().saturating_add(1));
2816 });
2817 EagerNoGradGuard {
2818 active: true,
2819 _not_send: PhantomData,
2820 }
2821 }
2822
2823 /// Keep semantic-trace recording active for untracked intermediates.
2824 ///
2825 /// See [`EagerTraceCaptureGuard`] for the full contract and an example.
2826 pub fn capture_trace(&self) -> EagerTraceCaptureGuard {
2827 EAGER_CAPTURE_DEPTH.with(|depth| {
2828 depth.set(depth.get().saturating_add(1));
2829 });
2830 EagerTraceCaptureGuard {
2831 active: true,
2832 _not_send: PhantomData,
2833 }
2834 }
2835
2836 /// Install or replace one extension module on this eager context's runtime.
2837 ///
2838 /// Eager extension wrappers call this as an idempotent "ensure installed"
2839 /// step. When the exact module instance (same module ID and allocation) is
2840 /// already installed, this is a read-only no-op that returns the current
2841 /// runtime epoch without acquiring the install lock or reconfiguring. The
2842 /// cold or replacement paths keep the transactional install-or-replace
2843 /// behavior, serialized so parallel first-use of the same extension family
2844 /// cannot publish over another thread's base snapshot.
2845 ///
2846 /// # Errors
2847 ///
2848 /// Returns [`tenferro_runtime::Error::RuntimeState`] when runtime
2849 /// reconfiguration fails or the extension module transaction is invalid.
2850 pub fn install_extension_module(
2851 &self,
2852 module: Arc<dyn ExtensionModule>,
2853 ) -> Result<RuntimeEpoch> {
2854 let snapshot = self.runtime.snapshot().map_err(|source| {
2855 runtime_state_source("EagerRuntime::install_extension_module", source)
2856 })?;
2857 if snapshot.has_extension_module_identical(&module) {
2858 return Ok(snapshot.epoch());
2859 }
2860 let _install_guard = self.lock_extension_install()?;
2861 self.runtime
2862 .reconfigure(|edit| {
2863 edit.replace_extension_module(module)?;
2864 Ok(())
2865 })
2866 .map_err(|source| {
2867 runtime_state_source("EagerRuntime::install_extension_module", source)
2868 })
2869 }
2870
2871 pub(crate) fn ensure_extension_module_for_engine(
2872 &self,
2873 module: Arc<dyn ExtensionModule>,
2874 family_id: &'static str,
2875 engine_id: &EngineId,
2876 ) -> Result<RuntimeEpoch> {
2877 let snapshot = self.runtime.snapshot().map_err(|source| {
2878 runtime_state_source("EagerRuntime::ensure_extension_module_for_engine", source)
2879 })?;
2880 if snapshot.has_extension_module_engine(module.module_id(), family_id, engine_id) {
2881 return Ok(snapshot.epoch());
2882 }
2883 let _install_guard = self.lock_extension_install()?;
2884 self.runtime
2885 .reconfigure(|edit| {
2886 edit.ensure_extension_module_for_engine(module, family_id, engine_id)?;
2887 Ok(())
2888 })
2889 .map_err(|source| {
2890 runtime_state_source("EagerRuntime::ensure_extension_module_for_engine", source)
2891 })
2892 }
2893
2894 pub(crate) fn runtime(&self) -> &Runtime {
2895 &self.runtime
2896 }
2897
2898 /// The extension target for `engine`, registering its kind first.
2899 pub(crate) fn eager_extension_target(
2900 &self,
2901 engine: &mut dyn ExecutionEngine,
2902 ) -> Result<EagerExtensionTarget> {
2903 let target = self.ensure_engine(engine)?;
2904 let backend_kind = eager_extension_backend_kind(&target).ok_or_else(|| {
2905 Error::unsupported(
2906 "EagerRuntime::eager_extension_target",
2907 ErrorPhase::Execution,
2908 "this engine kind has no eager extension target",
2909 )
2910 })?;
2911 let target = EagerExtensionTarget {
2912 engine_id: target.engine_id().clone(),
2913 backend_kind,
2914 };
2915 validate_eager_extension_target(&self.runtime, &target)?;
2916 Ok(target)
2917 }
2918
2919 /// Clear every cache owned by this eager context.
2920 ///
2921 /// # Examples
2922 ///
2923 /// ```
2924 /// use tenferro_ad::{EagerRuntime};
2925 ///
2926 /// let ctx = EagerRuntime::new()?;
2927 /// ctx.clear_caches()?;
2928 /// assert_eq!(ctx.cache_stats()?.ad_transforms.entries, 0);
2929 /// assert_eq!(ctx.cache_stats()?.prepared_derivatives.entries, 0);
2930 /// # Ok::<(), tenferro_ad::Error>(())
2931 /// ```
2932 ///
2933 /// # Errors
2934 ///
2935 /// Returns [`tenferro_runtime::Error::RuntimeState`] when either the
2936 /// extension cache or AD-transform cache is poisoned.
2937 pub fn clear_caches(&self) -> Result<()> {
2938 self.clear_ad_transform_caches()?;
2939 self.clear_prepared_derivative_cache()?;
2940 Ok(())
2941 }
2942
2943 /// Clear prepared derivative program cache entries.
2944 ///
2945 /// # Examples
2946 ///
2947 /// ```rust
2948 /// use tenferro_ad::EagerRuntime;
2949 ///
2950 /// let ctx = EagerRuntime::new()?;
2951 /// ctx.clear_prepared_derivative_cache()?;
2952 /// assert_eq!(ctx.cache_stats()?.prepared_derivatives.entries, 0);
2953 /// # Ok::<(), tenferro_ad::Error>(())
2954 /// ```
2955 ///
2956 /// # Errors
2957 ///
2958 /// Returns [`tenferro_runtime::Error::RuntimeState`] if the prepared
2959 /// derivative cache lock is poisoned.
2960 pub fn clear_prepared_derivative_cache(&self) -> Result<()> {
2961 self.lock_prepared_derivative_cache()?.clear();
2962 Ok(())
2963 }
2964
2965 /// Return eager runtime cache-entry and retained-byte stats.
2966 ///
2967 /// # Examples
2968 ///
2969 /// ```
2970 /// use tenferro_ad::{EagerRuntime};
2971 ///
2972 /// let ctx = EagerRuntime::new()?;
2973 /// let stats = ctx.cache_stats()?;
2974 /// assert_eq!(stats.ad_transforms.entries, 0);
2975 /// assert_eq!(stats.prepared_derivatives.entries, 0);
2976 /// # Ok::<(), tenferro_ad::Error>(())
2977 /// ```
2978 ///
2979 /// # Errors
2980 ///
2981 /// Returns [`tenferro_runtime::Error::RuntimeState`] when a cache or
2982 /// AD-transform cache lock is poisoned.
2983 pub fn cache_stats(&self) -> Result<EagerRuntimeCacheStats> {
2984 Ok(EagerRuntimeCacheStats {
2985 ad_transforms: self.ad_transform_cache.stats()?,
2986 prepared_derivatives: self.lock_prepared_derivative_cache()?.stats(),
2987 })
2988 }
2989
2990 /// Return the AD transform cache retention limits.
2991 ///
2992 /// # Examples
2993 ///
2994 /// ```
2995 /// use tenferro_ad::EagerRuntime;
2996 ///
2997 /// let ctx = EagerRuntime::new()?;
2998 /// assert!(ctx.ad_transform_cache_limits()?.max_entries().get() > 0);
2999 /// # Ok::<(), tenferro_ad::Error>(())
3000 /// ```
3001 ///
3002 /// # Errors
3003 ///
3004 /// Returns [`tenferro_runtime::Error::RuntimeState`] if the AD-transform
3005 /// cache lock is poisoned.
3006 pub fn ad_transform_cache_limits(&self) -> Result<AdTransformCacheLimits> {
3007 self.ad_transform_cache.limits()
3008 }
3009
3010 /// Replace AD transform cache retention limits.
3011 ///
3012 /// # Examples
3013 ///
3014 /// ```
3015 /// use std::num::NonZeroUsize;
3016 /// use tenferro_ad::{AdTransformCacheLimits, EagerRuntime};
3017 ///
3018 /// let ctx = EagerRuntime::new()?;
3019 /// let limits = AdTransformCacheLimits::new(NonZeroUsize::new(1).unwrap());
3020 /// ctx.set_ad_transform_cache_limits(limits)?;
3021 /// assert_eq!(ctx.ad_transform_cache_limits()?, limits);
3022 /// # Ok::<(), tenferro_ad::Error>(())
3023 /// ```
3024 ///
3025 /// # Errors
3026 ///
3027 /// Returns [`tenferro_runtime::Error::RuntimeState`] if the AD-transform
3028 /// cache lock is poisoned while updating limits.
3029 pub fn set_ad_transform_cache_limits(&self, limits: AdTransformCacheLimits) -> Result<()> {
3030 self.ad_transform_cache.set_limits(limits)
3031 }
3032
3033 /// Clear AD transform cache entries visible through this eager runtime.
3034 ///
3035 /// # Examples
3036 ///
3037 /// ```
3038 /// use tenferro_ad::EagerRuntime;
3039 ///
3040 /// let ctx = EagerRuntime::new()?;
3041 /// ctx.clear_ad_transform_caches()?;
3042 /// assert_eq!(ctx.cache_stats()?.ad_transforms.entries, 0);
3043 /// # Ok::<(), tenferro_ad::Error>(())
3044 /// ```
3045 ///
3046 /// # Errors
3047 ///
3048 /// Returns [`tenferro_runtime::Error::RuntimeState`] if the AD-transform
3049 /// cache lock is poisoned while clearing entries.
3050 pub fn clear_ad_transform_caches(&self) -> Result<()> {
3051 self.ad_transform_cache.clear()
3052 }
3053
3054 /// Return prepared derivative cache retention limits.
3055 ///
3056 /// # Examples
3057 ///
3058 /// ```rust
3059 /// use tenferro_ad::EagerRuntime;
3060 ///
3061 /// let ctx = EagerRuntime::new()?;
3062 /// assert!(ctx.prepared_derivative_cache_limits()?.max_entries().get() > 0);
3063 /// # Ok::<(), tenferro_ad::Error>(())
3064 /// ```
3065 ///
3066 /// # Errors
3067 ///
3068 /// Returns [`tenferro_runtime::Error::RuntimeState`] if the prepared
3069 /// derivative cache lock is poisoned.
3070 pub fn prepared_derivative_cache_limits(&self) -> Result<AdTransformCacheLimits> {
3071 Ok(self.lock_prepared_derivative_cache()?.limits())
3072 }
3073
3074 /// Replace prepared derivative cache retention limits.
3075 ///
3076 /// # Examples
3077 ///
3078 /// ```rust
3079 /// use std::num::NonZeroUsize;
3080 /// use tenferro_ad::{AdTransformCacheLimits, EagerRuntime};
3081 ///
3082 /// let ctx = EagerRuntime::new()?;
3083 /// let limits = AdTransformCacheLimits::new(NonZeroUsize::new(1).unwrap());
3084 /// ctx.set_prepared_derivative_cache_limits(limits)?;
3085 /// assert_eq!(ctx.prepared_derivative_cache_limits()?, limits);
3086 /// # Ok::<(), tenferro_ad::Error>(())
3087 /// ```
3088 ///
3089 /// # Errors
3090 ///
3091 /// Returns [`tenferro_runtime::Error::RuntimeState`] if the prepared
3092 /// derivative cache lock is poisoned.
3093 pub fn set_prepared_derivative_cache_limits(
3094 &self,
3095 limits: AdTransformCacheLimits,
3096 ) -> Result<()> {
3097 self.lock_prepared_derivative_cache()?.set_limits(limits);
3098 Ok(())
3099 }
3100
3101 /// Run eager operations on `engine`.
3102 ///
3103 /// The callback receives an [`EagerSession`] that pairs this runtime with
3104 /// the borrowed engine; operations recorded through it are differentiable.
3105 /// Nothing is entered or locked: the engine is simply borrowed for the
3106 /// callback.
3107 ///
3108 /// # Examples
3109 ///
3110 /// ```
3111 /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
3112 /// use tenferro_cpu::CpuEngine;
3113 ///
3114 /// let ctx = EagerRuntime::new()?;
3115 /// let mut engine = CpuEngine::serial();
3116 /// let x = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![1], vec![1.0_f64])?, ctx.clone())?;
3117 /// let y = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![1], vec![2.0_f64])?, ctx.clone())?;
3118 /// let z = ctx.with_eager_session(&mut engine, |s| s.add(&x, &y))?;
3119 /// assert_eq!(z.value()?.as_slice::<f64>()?, &[3.0]);
3120 /// # Ok::<(), tenferro_ad::Error>(())
3121 /// ```
3122 ///
3123 /// # Errors
3124 ///
3125 /// Returns [`Error::Unsupported`] when the engine kind cannot be registered
3126 /// with this runtime or its context identity conflicts with the existing
3127 /// registration, [`Error::RuntimeStateSource`] when the runtime snapshot or
3128 /// engine registration state is unavailable, and the callback's own error
3129 /// unchanged.
3130 pub fn with_eager_session<T, E: From<Error>>(
3131 self: &Arc<Self>,
3132 engine: &mut dyn ExecutionEngine,
3133 f: impl FnOnce(&mut EagerSession<'_>) -> std::result::Result<T, E>,
3134 ) -> std::result::Result<T, E> {
3135 self.ensure_engine(engine).map_err(E::from)?;
3136 f(&mut EagerSession {
3137 runtime: self,
3138 engine,
3139 })
3140 }
3141
3142 /// Run a compiled program on `engine` through this runtime.
3143 pub(crate) fn run_program(
3144 &self,
3145 engine: &mut dyn ExecutionEngine,
3146 program: &CompiledGraph,
3147 inputs: &[TensorRead<'_>],
3148 ) -> Result<Vec<Tensor>> {
3149 let target = self.ensure_engine(engine)?;
3150 let prepared = self
3151 .runtime
3152 .prepare(program, inputs, std::slice::from_ref(&target))?;
3153 run_prepared_on(&self.runtime, engine, &prepared, inputs)
3154 }
3155
3156 pub(crate) fn exec_extension_outputs_read(
3157 &self,
3158 engine: &mut dyn ExecutionEngine,
3159 op: &Arc<dyn tenferro_ops::ext_op::ExtensionOp>,
3160 inputs: &[TensorRead<'_>],
3161 ) -> Result<Vec<Tensor>> {
3162 self.ensure_engine(engine)?;
3163 profile_eager_op_section("exec_extension_outputs_read.exec_op", || {
3164 exec_extension_op_on_tensor_reads(op, inputs, engine, &self.runtime)
3165 })
3166 }
3167
3168 #[cfg(test)]
3169 pub(crate) fn exec_standard_graph_outputs(
3170 &self,
3171 engine: &mut dyn ExecutionEngine,
3172 graph: &Graph<StdTensorOp>,
3173 initial_data: HashMap<ValueKey<StdTensorOp>, Tensor>,
3174 ) -> Result<EagerGraphExecution> {
3175 let mut all_values = initial_data;
3176
3177 profile_eager_op_section("exec_graph.run", || {
3178 let exec: &mut dyn Engine = engine;
3179 (|| -> Result<()> {
3180 for op_node in graph.operations() {
3181 let outputs = {
3182 let input_values = op_node
3183 .inputs
3184 .iter()
3185 .map(|input| {
3186 let key = match input {
3187 ValueRef::Local(local_id) => &graph.values()[*local_id].key,
3188 ValueRef::External(key) => key,
3189 };
3190 all_values.get(key).ok_or_else(|| {
3191 Error::Internal(format!(
3192 "standard graph eager execution missing value for {key:?}"
3193 ))
3194 })
3195 })
3196 .collect::<Result<Vec<_>>>()?;
3197 let input_reads = input_values
3198 .iter()
3199 .map(|value| TensorRead::from_tensor(value))
3200 .collect::<Vec<_>>();
3201 exec_standard_op_on_tensor_reads_in_session(
3202 &op_node.operation,
3203 &input_reads,
3204 exec,
3205 )?
3206 };
3207
3208 if outputs.len() != op_node.outputs.len() {
3209 return Err(Error::Internal(format!(
3210 "standard graph eager execution expected {} outputs for {:?}, got {}",
3211 op_node.outputs.len(),
3212 op_node.operation,
3213 outputs.len()
3214 )));
3215 }
3216
3217 for (output_id, output) in op_node.outputs.iter().zip(outputs) {
3218 let key = graph.values()[*output_id].key.clone();
3219 all_values.insert(key, output);
3220 }
3221 }
3222 Ok(())
3223 })()
3224 })?;
3225
3226 let outputs = graph
3227 .outputs()
3228 .iter()
3229 .map(|&output_id| {
3230 let key = &graph.values()[output_id].key;
3231 all_values
3232 .get(key)
3233 .ok_or_else(|| {
3234 Error::Internal(format!(
3235 "standard graph eager execution missing graph output {key:?}"
3236 ))
3237 })?
3238 .duplicate()
3239 .map_err(Error::from)
3240 })
3241 .collect::<Result<Vec<_>>>()?;
3242
3243 Ok(EagerGraphExecution { outputs })
3244 }
3245
3246 pub(crate) fn try_register_grad_slot(
3247 &self,
3248 key: &ValueKey<StdTensorOp>,
3249 slot: &GradSlot,
3250 ) -> Result<()> {
3251 insert_pruning_dead(
3252 &mut *self.lock_grad_slots()?,
3253 key.clone(),
3254 Arc::downgrade(slot),
3255 );
3256 Ok(())
3257 }
3258
3259 pub(crate) fn try_register_value_record(
3260 &self,
3261 key: &ValueKey<StdTensorOp>,
3262 record: &Arc<EagerTensorRecord>,
3263 ) -> Result<()> {
3264 insert_pruning_dead(
3265 &mut *self.lock_value_records()?,
3266 key.clone(),
3267 Arc::downgrade(record),
3268 );
3269 Ok(())
3270 }
3271
3272 pub(crate) fn value_record(
3273 &self,
3274 key: &ValueKey<StdTensorOp>,
3275 ) -> Result<Option<Arc<EagerTensorRecord>>> {
3276 let mut records = self.lock_value_records()?;
3277 let Some(record) = records.get(key).cloned() else {
3278 return Ok(None);
3279 };
3280 match record.upgrade() {
3281 Some(record) => Ok(Some(record)),
3282 None => {
3283 records.remove(key);
3284 Ok(None)
3285 }
3286 }
3287 }
3288
3289 /// Clear all live gradient slots tracked by this context.
3290 ///
3291 /// This resets the stored gradients to `None` without unregistering the
3292 /// tensors, so future `backward()` calls can accumulate again.
3293 ///
3294 /// # Examples
3295 ///
3296 /// ```
3297 /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
3298 ///
3299 /// let ctx = EagerRuntime::new()?;
3300 /// let x = EagerTensor::requires_grad_in(Tensor::from_vec_col_major(vec![3], vec![1.0_f64, 2.0, 3.0]).unwrap(), ctx.clone()).unwrap();
3301 /// let y = EagerTensor::requires_grad_in(Tensor::from_vec_col_major(vec![3], vec![4.0_f64, 5.0, 6.0]).unwrap(), ctx.clone()).unwrap();
3302 /// let loss = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
3303 /// let product = s.mul(&x, &y)?;
3304 /// s.reduce_sum(&product, Some(&[0]))
3305 /// })?;
3306 /// let _ = loss.backward(&mut tenferro_cpu::CpuEngine::serial()).unwrap();
3307 ///
3308 /// ctx.clear_grads()?;
3309 ///
3310 /// assert!(x.grad()?.is_none());
3311 /// assert!(y.grad()?.is_none());
3312 /// # Ok::<(), tenferro_ad::Error>(())
3313 /// ```
3314 ///
3315 /// # Errors
3316 ///
3317 /// Returns [`tenferro_runtime::Error::RuntimeState`] if a gradient-slot
3318 /// lock is poisoned while clearing live gradients.
3319 pub fn clear_grads(&self) -> Result<()> {
3320 let live_slots = {
3321 let mut live_slots = Vec::new();
3322 self.lock_grad_slots()?.retain(|_, slot| {
3323 if let Some(slot) = slot.upgrade() {
3324 live_slots.push(slot);
3325 true
3326 } else {
3327 false
3328 }
3329 });
3330 live_slots
3331 };
3332
3333 let mut poisoned_slot = false;
3334 for slot in live_slots {
3335 match slot.lock() {
3336 Ok(mut current) => {
3337 *current = None;
3338 }
3339 Err(_) => {
3340 poisoned_slot = true;
3341 }
3342 }
3343 }
3344 if poisoned_slot {
3345 return Err(Error::runtime_state(
3346 "eager_gradient_slot",
3347 ErrorPhase::Execution,
3348 "lock poisoned",
3349 ));
3350 }
3351 Ok(())
3352 }
3353
3354 /// Import a concrete tensor into this context as an untracked constant.
3355 ///
3356 /// The returned tensor does not participate in gradient tracking.
3357 /// Use this for fixed masks, quadrature weights, physical constants,
3358 /// and other data that should not receive gradients. Like
3359 /// [`EagerSession::constant_from`], it performs no host/device transfer;
3360 /// use [`EagerSession::constant_from_host`] to upload host data into a
3361 /// device runtime.
3362 ///
3363 /// # Examples
3364 ///
3365 /// ```
3366 /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
3367 ///
3368 /// let ctx = EagerRuntime::new()?;
3369 /// let c = ctx.constant_from(Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0]).unwrap())?;
3370 /// let x = EagerTensor::requires_grad_in(Tensor::from_vec_col_major(vec![2], vec![3.0_f64, 4.0]).unwrap(), ctx.clone())?;
3371 /// let z = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| s.add(&x, &c))?;
3372 ///
3373 /// assert_eq!(z.value()?.as_slice::<f64>().unwrap(), &[4.0, 6.0]);
3374 /// # Ok::<(), tenferro_ad::Error>(())
3375 /// ```
3376 ///
3377 /// # Errors
3378 ///
3379 /// Returns [`tenferro_runtime::Error::RuntimeState`] when metadata cannot
3380 /// be registered or the backend lock is poisoned.
3381 pub fn constant_from(self: &Arc<Self>, tensor: Tensor) -> Result<EagerTensor> {
3382 EagerTensor::new_leaf(Arc::clone(self), tensor, false)
3383 }
3384
3385 /// Import a concrete tensor into this context as a trainable variable.
3386 ///
3387 /// The returned tensor participates in gradient tracking; its gradient
3388 /// slot is registered in this context.
3389 ///
3390 /// # Examples
3391 ///
3392 /// ```
3393 /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
3394 ///
3395 /// let ctx = EagerRuntime::new()?;
3396 /// let p = ctx.variable_from(Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0]).unwrap())?;
3397 /// let loss = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
3398 /// let y = s.exp(&p)?;
3399 /// s.reduce_sum(&y, Some(&[0]))
3400 /// })?;
3401 /// let _ = loss.backward(&mut tenferro_cpu::CpuEngine::serial()).unwrap();
3402 ///
3403 /// let grad = p.grad().unwrap().unwrap();
3404 /// assert_eq!(grad.shape(), &[2]);
3405 /// # Ok::<(), tenferro_ad::Error>(())
3406 /// ```
3407 ///
3408 /// # Errors
3409 ///
3410 /// Returns [`tenferro_runtime::Error::RuntimeState`] when gradient metadata
3411 /// or the eager backend state cannot be registered.
3412 pub fn variable_from(self: &Arc<Self>, tensor: Tensor) -> Result<EagerTensor> {
3413 EagerTensor::new_leaf(Arc::clone(self), tensor, true)
3414 }
3415
3416 /// Gradient of a scalar eager output with respect to an eager tensor.
3417 ///
3418 /// Functional eager gradients return ordinary eager tensors and do not
3419 /// write into `grad()` slots. The returned tensor keeps a trace when the
3420 /// derivative computation depends on tracked eager values.
3421 ///
3422 /// # Examples
3423 ///
3424 /// ```
3425 /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
3426 ///
3427 /// let ctx = EagerRuntime::new()?;
3428 /// let x = EagerTensor::requires_grad_in(
3429 /// Tensor::from_vec_col_major(vec![], vec![3.0_f64]).unwrap(),
3430 /// ctx.clone(),
3431 /// )?;
3432 /// let loss = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| s.mul(&x, &x))?;
3433 /// let dx = ctx.grad(&mut tenferro_cpu::CpuEngine::serial(), &loss, &x)?;
3434 /// assert_eq!(dx.value()?.as_slice::<f64>().unwrap(), &[6.0]);
3435 /// # Ok::<(), tenferro_ad::Error>(())
3436 /// ```
3437 ///
3438 /// # Errors
3439 ///
3440 /// Returns [`tenferro_runtime::Error::NonScalarGrad`] for a non-scalar
3441 /// output, [`Error::ContextMismatch`] for tensors from another runtime,
3442 /// [`Error::UnsupportedAdRule`] when an AD rule is unavailable, or a typed
3443 /// validation/backend error from eager execution. An inactive `wrt` returns
3444 /// [`Error::Validation`] with `argument: "wrt"`; use
3445 /// [`grad_optional`](Self::grad_optional) to observe that state.
3446 pub fn grad(
3447 self: &Arc<Self>,
3448 engine: &mut dyn ExecutionEngine,
3449 output: &EagerTensor,
3450 wrt: &EagerTensor,
3451 ) -> Result<EagerTensor> {
3452 self.grad_optional(engine, output, wrt)?
3453 .ok_or_else(|| crate::traced::inactive_wrt_error("grad", &wrt.key))
3454 }
3455
3456 /// Gradient that returns `None` when `wrt` is inactive.
3457 ///
3458 /// # Examples
3459 ///
3460 /// ```
3461 /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
3462 ///
3463 /// let ctx = EagerRuntime::new()?;
3464 /// let x = EagerTensor::requires_grad_in(
3465 /// Tensor::from_vec_col_major(vec![], vec![3.0_f64]).unwrap(),
3466 /// ctx.clone(),
3467 /// )?;
3468 /// let y = EagerTensor::requires_grad_in(
3469 /// Tensor::from_vec_col_major(vec![], vec![4.0_f64]).unwrap(),
3470 /// ctx.clone(),
3471 /// )?;
3472 /// let loss = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| s.mul(&y, &y))?;
3473 /// assert!(ctx.grad_optional(&mut tenferro_cpu::CpuEngine::serial(), &loss, &x)?.is_none());
3474 /// # Ok::<(), tenferro_ad::Error>(())
3475 /// ```
3476 ///
3477 /// # Errors
3478 ///
3479 /// Returns [`tenferro_runtime::Error::NonScalarGrad`] for a non-scalar
3480 /// output, [`Error::ContextMismatch`] for a foreign runtime, or a typed
3481 /// validation/backend/runtime-state error from eager execution.
3482 pub fn grad_optional(
3483 self: &Arc<Self>,
3484 engine: &mut dyn ExecutionEngine,
3485 output: &EagerTensor,
3486 wrt: &EagerTensor,
3487 ) -> Result<Option<EagerTensor>> {
3488 if !output.shape().is_empty() {
3489 return Err(Error::NonScalarGrad {
3490 shape: output.shape().to_vec(),
3491 });
3492 }
3493
3494 let value = output.to_tensor(&mut *engine)?;
3495 let seed = one_like_tensor(&value, &mut *engine)?;
3496 let seed = EagerTensor::new_result(Arc::clone(self), eager_val_key(), seed, false, None)?;
3497 self.vjp_optional(engine, output, wrt, &seed)
3498 }
3499
3500 /// Reverse-mode vector-Jacobian product for eager tensors.
3501 ///
3502 /// # Examples
3503 ///
3504 /// ```
3505 /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
3506 ///
3507 /// let ctx = EagerRuntime::new()?;
3508 /// let x = EagerTensor::requires_grad_in(
3509 /// Tensor::from_vec_col_major(vec![2], vec![2.0_f64, 3.0]).unwrap(),
3510 /// ctx.clone(),
3511 /// )?;
3512 /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| s.mul(&x, &x))?;
3513 /// let seed = EagerTensor::from_tensor_in(
3514 /// Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 1.0]).unwrap(),
3515 /// ctx.clone(),
3516 /// )?;
3517 /// let dx = ctx.vjp(&mut tenferro_cpu::CpuEngine::serial(), &y, &x, &seed)?;
3518 /// assert_eq!(dx.value()?.as_slice::<f64>().unwrap(), &[4.0, 6.0]);
3519 /// # Ok::<(), tenferro_ad::Error>(())
3520 /// ```
3521 ///
3522 /// # Errors
3523 ///
3524 /// Returns [`Error::ContextMismatch`] for tensors from different eager
3525 /// runtimes, [`Error::Validation`] when the cotangent shape or dtype does
3526 /// not match the output, [`Error::UnsupportedAdRule`] when a rule is not
3527 /// registered, or a typed backend/runtime-state error. An inactive `wrt`
3528 /// returns [`Error::Validation`] with `argument: "wrt"`; use
3529 /// [`vjp_optional`](Self::vjp_optional) to observe that state.
3530 pub fn vjp(
3531 self: &Arc<Self>,
3532 engine: &mut dyn ExecutionEngine,
3533 output: &EagerTensor,
3534 wrt: &EagerTensor,
3535 cotangent: &EagerTensor,
3536 ) -> Result<EagerTensor> {
3537 self.vjp_optional(engine, output, wrt, cotangent)?
3538 .ok_or_else(|| crate::traced::inactive_wrt_error("vjp", &wrt.key))
3539 }
3540
3541 /// Reverse-mode vector-Jacobian product that returns `None` for inactive inputs.
3542 ///
3543 /// `None` means the input is disconnected, its eager computation was not
3544 /// recorded, or the semantic derivative has no cotangent flow (for example,
3545 /// real `sign` or an integer conversion). A present numerical zero is still
3546 /// returned as `Some`. Missing rules and execution failures return errors;
3547 /// semantic differentiation cannot be disabled by an environment variable.
3548 ///
3549 /// # Examples
3550 ///
3551 /// ```
3552 /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
3553 ///
3554 /// let ctx = EagerRuntime::new()?;
3555 /// let x = EagerTensor::requires_grad_in(
3556 /// Tensor::from_vec_col_major(vec![1], vec![2.0_f64]).unwrap(),
3557 /// ctx.clone(),
3558 /// )?;
3559 /// let y = EagerTensor::requires_grad_in(
3560 /// Tensor::from_vec_col_major(vec![1], vec![4.0_f64]).unwrap(),
3561 /// ctx.clone(),
3562 /// )?;
3563 /// let seed = EagerTensor::from_tensor_in(
3564 /// Tensor::from_vec_col_major(vec![1], vec![1.0_f64]).unwrap(),
3565 /// ctx.clone(),
3566 /// )?;
3567 /// let loss = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| s.mul(&y, &y))?;
3568 /// assert!(ctx.vjp_optional(&mut tenferro_cpu::CpuEngine::serial(), &loss, &x, &seed)?.is_none());
3569 /// # Ok::<(), tenferro_ad::Error>(())
3570 /// ```
3571 ///
3572 /// # Errors
3573 ///
3574 /// Returns [`Error::ContextMismatch`] for tensors from different eager
3575 /// runtimes, [`Error::Validation`] when the cotangent shape or dtype does
3576 /// not match the output, [`Error::UnsupportedAdRule`] when a rule is not
3577 /// registered, or a typed backend/runtime-state error.
3578 pub fn vjp_optional(
3579 self: &Arc<Self>,
3580 engine: &mut dyn ExecutionEngine,
3581 output: &EagerTensor,
3582 wrt: &EagerTensor,
3583 cotangent: &EagerTensor,
3584 ) -> Result<Option<EagerTensor>> {
3585 validate_same_runtime(self, output, "vjp output")?;
3586 validate_same_runtime(self, wrt, "vjp wrt")?;
3587 validate_same_runtime(self, cotangent, "vjp cotangent")?;
3588 validate_seed_tensor("vjp", output, cotangent)?;
3589 Ok(
3590 semantic_eager_vjp_many(self, engine, output, &[wrt], cotangent)?
3591 .pop()
3592 .flatten(),
3593 )
3594 }
3595
3596 /// Forward-mode Jacobian-vector product for eager tensors.
3597 ///
3598 /// # Examples
3599 ///
3600 /// ```
3601 /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
3602 ///
3603 /// let ctx = EagerRuntime::new()?;
3604 /// let x = EagerTensor::requires_grad_in(
3605 /// Tensor::from_vec_col_major(vec![1], vec![3.0_f64]).unwrap(),
3606 /// ctx.clone(),
3607 /// )?;
3608 /// let tangent = EagerTensor::from_tensor_in(
3609 /// Tensor::from_vec_col_major(vec![1], vec![1.0_f64]).unwrap(),
3610 /// ctx.clone(),
3611 /// )?;
3612 /// let y = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| s.mul(&x, &x))?;
3613 /// let dy = ctx.jvp(&mut tenferro_cpu::CpuEngine::serial(), &y, &x, &tangent)?;
3614 /// assert_eq!(dy.value()?.as_slice::<f64>().unwrap(), &[6.0]);
3615 /// # Ok::<(), tenferro_ad::Error>(())
3616 /// ```
3617 ///
3618 /// # Errors
3619 ///
3620 /// Returns [`Error::ContextMismatch`] for tensors from different eager
3621 /// runtimes, [`Error::Validation`] when the tangent shape or dtype does not
3622 /// match `wrt`, [`Error::UnsupportedAdRule`] when a rule is unavailable, or
3623 /// a typed backend/runtime-state error. An inactive `wrt` returns
3624 /// [`Error::Validation`] with `argument: "wrt"`; use
3625 /// [`jvp_optional`](Self::jvp_optional) to observe that state.
3626 pub fn jvp(
3627 self: &Arc<Self>,
3628 engine: &mut dyn ExecutionEngine,
3629 output: &EagerTensor,
3630 wrt: &EagerTensor,
3631 tangent: &EagerTensor,
3632 ) -> Result<EagerTensor> {
3633 self.jvp_optional(engine, output, wrt, tangent)?
3634 .ok_or_else(|| crate::traced::inactive_wrt_error("jvp", &wrt.key))
3635 }
3636
3637 /// Forward-mode Jacobian-vector product that returns `None` for inactive outputs.
3638 ///
3639 /// `None` means `wrt` is disconnected, its eager computation was not
3640 /// recorded, or the semantic derivative has no tangent flow (for example,
3641 /// real `sign` or an integer conversion). A present numerical zero is still
3642 /// returned as `Some`. Missing rules and execution failures return errors;
3643 /// semantic differentiation cannot be disabled by an environment variable.
3644 ///
3645 /// # Examples
3646 ///
3647 /// ```
3648 /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
3649 ///
3650 /// let ctx = EagerRuntime::new()?;
3651 /// let x = EagerTensor::requires_grad_in(
3652 /// Tensor::from_vec_col_major(vec![1], vec![2.0_f64]).unwrap(),
3653 /// ctx.clone(),
3654 /// )?;
3655 /// let y = EagerTensor::requires_grad_in(
3656 /// Tensor::from_vec_col_major(vec![1], vec![4.0_f64]).unwrap(),
3657 /// ctx.clone(),
3658 /// )?;
3659 /// let tangent = EagerTensor::from_tensor_in(
3660 /// Tensor::from_vec_col_major(vec![1], vec![1.0_f64]).unwrap(),
3661 /// ctx.clone(),
3662 /// )?;
3663 /// let loss = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| s.mul(&y, &y))?;
3664 /// assert!(ctx.jvp_optional(&mut tenferro_cpu::CpuEngine::serial(), &loss, &x, &tangent)?.is_none());
3665 /// # Ok::<(), tenferro_ad::Error>(())
3666 /// ```
3667 ///
3668 /// # Errors
3669 ///
3670 /// Returns [`Error::ContextMismatch`] for tensors from different eager
3671 /// runtimes, [`Error::Validation`] when the tangent shape or dtype does not
3672 /// match `wrt`, [`Error::UnsupportedAdRule`] when a rule is unavailable, or
3673 /// a typed backend/runtime-state error.
3674 pub fn jvp_optional(
3675 self: &Arc<Self>,
3676 engine: &mut dyn ExecutionEngine,
3677 output: &EagerTensor,
3678 wrt: &EagerTensor,
3679 tangent: &EagerTensor,
3680 ) -> Result<Option<EagerTensor>> {
3681 validate_same_runtime(self, output, "jvp output")?;
3682 validate_same_runtime(self, wrt, "jvp wrt")?;
3683 validate_same_runtime(self, tangent, "jvp tangent")?;
3684 validate_seed_tensor("jvp", wrt, tangent)?;
3685 semantic_eager_jvp_optional(self, engine, output, wrt, tangent)
3686 }
3687
3688 fn store_grads(
3689 &self,
3690 cotangents: &HashMap<ValueKey<StdTensorOp>, Tensor>,
3691 session: &mut dyn Engine,
3692 ) -> Result<()> {
3693 let mut updates = Vec::new();
3694
3695 {
3696 let mut slots = self.lock_grad_slots()?;
3697 slots.retain(|key, slot| {
3698 let Some(slot) = slot.upgrade() else {
3699 return false;
3700 };
3701
3702 if let Some(incoming) = cotangents.get(key) {
3703 updates.push((slot, incoming));
3704 }
3705
3706 true
3707 });
3708 }
3709
3710 for (slot, incoming) in updates {
3711 let mut current = slot.lock().map_err(|_| {
3712 Error::runtime_state(
3713 "eager_gradient_slot",
3714 ErrorPhase::Execution,
3715 "lock poisoned",
3716 )
3717 })?;
3718 let next = match current.as_ref() {
3719 Some(existing) => {
3720 let existing_read = existing.tensor_read("EagerRuntime::store_grads")?;
3721 let incoming_read = TensorRead::from_tensor(incoming);
3722 let tensor = session
3723 .add_read(existing_read, incoming_read)
3724 .map_err(Error::from)?;
3725 AdValueRecord::from_tensor(tensor, "EagerRuntime::store_grads")?
3726 }
3727 None => {
3728 let duplicate = session
3729 .to_contiguous_read(TensorRead::from_tensor(incoming))
3730 .map_err(Error::from)?;
3731 AdValueRecord::from_tensor(duplicate, "EagerRuntime::store_grads")?
3732 }
3733 };
3734 *current = Some(next);
3735 }
3736
3737 Ok(())
3738 }
3739}
3740
3741#[derive(Clone, Debug, PartialEq, Eq, Hash)]
3742struct PreparedDerivativeCacheKey {
3743 /// Engines with the same target share prepared derivative programs.
3744 target: EngineTarget,
3745 semantic_fingerprint: SemanticFingerprint,
3746 runtime_epoch: RuntimeEpoch,
3747 active_inputs: Box<[bool]>,
3748 input_metadata: Box<[ProgramValueMetadata]>,
3749}
3750
3751/// Cached prepared derivative: program + index metadata.
3752#[derive(Debug)]
3753struct PreparedDerivative {
3754 program: Arc<CompiledGraph>,
3755 execution_program: Arc<CompiledGraph>,
3756 saved_input_indices: Vec<usize>,
3757 prepared: Arc<Prepared>,
3758 seed_input_index: usize,
3759 derivative_output_indices: Box<[Option<usize>]>,
3760}
3761
3762#[derive(Debug)]
3763struct PreparedDerivativeCache {
3764 limits: AdTransformCacheLimits,
3765 entries: LruCache<PreparedDerivativeCacheKey, PreparedDerivativeCacheEntry>,
3766 stats: CacheStats,
3767}
3768
3769impl PreparedDerivativeCache {
3770 fn limits(&self) -> AdTransformCacheLimits {
3771 self.limits
3772 }
3773
3774 fn set_limits(&mut self, limits: AdTransformCacheLimits) {
3775 self.limits = limits;
3776 self.evict_to_limits();
3777 }
3778
3779 fn clear(&mut self) {
3780 let clears = self.stats.clears.saturating_add(1);
3781 self.entries.clear();
3782 self.stats = CacheStats {
3783 clears,
3784 ..CacheStats::empty()
3785 };
3786 }
3787
3788 fn stats(&self) -> CacheStats {
3789 self.stats
3790 }
3791
3792 fn get(&mut self, key: &PreparedDerivativeCacheKey) -> Option<Arc<PreparedDerivative>> {
3793 match self.entries.get(key) {
3794 Some(entry) => {
3795 self.stats.hits = self.stats.hits.saturating_add(1);
3796 Some(Arc::clone(&entry.value))
3797 }
3798 None => {
3799 self.stats.misses = self.stats.misses.saturating_add(1);
3800 None
3801 }
3802 }
3803 }
3804
3805 fn insert(&mut self, key: PreparedDerivativeCacheKey, value: Arc<PreparedDerivative>) {
3806 let retained_bytes = prepared_derivative_cache_entry_retained_bytes(&key, value.as_ref());
3807 let entry = PreparedDerivativeCacheEntry {
3808 value,
3809 retained_bytes,
3810 };
3811 self.stats.retained_bytes = self.stats.retained_bytes.saturating_add(retained_bytes);
3812 if let Some((_old_key, old_entry)) = self.entries.push(key, entry) {
3813 self.stats.retained_bytes = self
3814 .stats
3815 .retained_bytes
3816 .saturating_sub(old_entry.retained_bytes);
3817 }
3818 self.stats.entries = self.entries.len();
3819 self.evict_to_limits();
3820 }
3821
3822 fn evict_to_limits(&mut self) {
3823 while self.entries.len() > self.limits.max_entries().get()
3824 || self
3825 .limits
3826 .max_retained_bytes()
3827 .is_some_and(|limit| self.stats.retained_bytes > limit.get())
3828 {
3829 let Some((_key, entry)) = self.entries.pop_lru() else {
3830 break;
3831 };
3832 self.stats.retained_bytes = self
3833 .stats
3834 .retained_bytes
3835 .saturating_sub(entry.retained_bytes);
3836 self.stats.evictions = self.stats.evictions.saturating_add(1);
3837 }
3838 self.stats.entries = self.entries.len();
3839 }
3840}
3841
3842impl Default for PreparedDerivativeCache {
3843 fn default() -> Self {
3844 Self {
3845 limits: AdTransformCacheLimits::default(),
3846 entries: LruCache::unbounded(),
3847 stats: CacheStats::empty(),
3848 }
3849 }
3850}
3851
3852#[derive(Debug)]
3853struct PreparedDerivativeCacheEntry {
3854 value: Arc<PreparedDerivative>,
3855 retained_bytes: usize,
3856}
3857
3858fn prepared_derivative_cache_entry_retained_bytes(
3859 key: &PreparedDerivativeCacheKey,
3860 value: &PreparedDerivative,
3861) -> usize {
3862 size_of::<PreparedDerivativeCacheKey>()
3863 .saturating_add(size_of_val(key.active_inputs.as_ref()))
3864 .saturating_add(size_of_val(value.derivative_output_indices.as_ref()))
3865 .saturating_add(size_of_val(value.saved_input_indices.as_slice()))
3866 .saturating_add(compiled_graph_retained_bytes(
3867 value.execution_program.as_ref(),
3868 ))
3869 .saturating_add(
3870 key.input_metadata
3871 .len()
3872 .saturating_mul(size_of::<ProgramValueMetadata>()),
3873 )
3874 .saturating_add(size_of::<PreparedDerivative>())
3875 .saturating_add(compiled_graph_retained_bytes(value.program.as_ref()))
3876 .saturating_add(prepared_compiled_graph_retained_bytes(
3877 value.prepared.as_ref(),
3878 value.program.as_ref(),
3879 ))
3880}
3881
3882fn prepared_compiled_graph_retained_bytes(
3883 prepared: &Prepared,
3884 derivative_program: &CompiledGraph,
3885) -> usize {
3886 size_of_val(prepared).saturating_add(compiled_graph_retained_bytes(derivative_program))
3887}
3888
3889fn compiled_graph_retained_bytes(program: &CompiledGraph) -> usize {
3890 size_of::<CompiledGraph>()
3891 .saturating_add(size_of_val(program.input_keys()))
3892 .saturating_add(program.bindings().len().saturating_mul(size_of::<usize>()))
3893 .saturating_add(semantic_program_retained_bytes(program.program()))
3894}
3895
3896fn semantic_program_retained_bytes(program: &SemanticProgram) -> usize {
3897 size_of::<SemanticProgram>()
3898 .saturating_add(size_of_val(program.inputs()))
3899 .saturating_add(size_of_val(program.outputs()))
3900 .saturating_add(
3901 program
3902 .operations()
3903 .len()
3904 .saturating_mul(size_of::<usize>()),
3905 )
3906 .saturating_add(
3907 program
3908 .shape_guards()
3909 .len()
3910 .saturating_mul(size_of::<usize>()),
3911 )
3912}
3913
3914fn semantic_eager_vjp_many(
3915 ctx: &Arc<EagerRuntime>,
3916 engine: &mut dyn ExecutionEngine,
3917 output: &EagerTensor,
3918 wrts: &[&EagerTensor],
3919 cotangent: &EagerTensor,
3920) -> Result<Vec<Option<EagerTensor>>> {
3921 if wrts.is_empty() {
3922 return Ok(Vec::new());
3923 }
3924 let Some(raw_output_trace) = output.semantic_trace.as_ref() else {
3925 return Ok(vec![None; wrts.len()]);
3926 };
3927 if !wrts.iter().any(|wrt| {
3928 wrt.semantic_trace
3929 .as_ref()
3930 .and_then(TracedTensor::input_key)
3931 .is_some_and(|key| raw_output_trace.has_attached_input_key(&key))
3932 }) {
3933 return Ok(vec![None; wrts.len()]);
3934 }
3935
3936 // First AD request on this output: run the deferred graph analysis over
3937 // the whole raw carrier chain once (metadata registration + constraint
3938 // scopes), so `compile_ad_source` sees the same analyzed graph the eager
3939 // forward used to append.
3940 let output_trace = analyze_deferred_semantic_trace(raw_output_trace)?;
3941 let saved = output
3942 .trace
3943 .as_ref()
3944 .map(EagerTrace::collect)
3945 .unwrap_or_default();
3946 let mut source_outputs = vec![&output_trace];
3947 source_outputs.extend(saved.iter().map(|value| &value.trace));
3948
3949 // Residual roots keep the semantic producers available to differentiation;
3950 // only the execution-only derivative replaces their numerical values.
3951 let mut compiler = GraphCompiler::new();
3952 let source =
3953 tenferro_runtime::ad_support::compile_ad_source_many(&mut compiler, &source_outputs)?;
3954 if source.output_count() != source_outputs.len()
3955 || source.input_keys().len() != source.input_count()
3956 || source.bindings().len() != source.input_count()
3957 {
3958 return Err(Error::Internal(
3959 "semantic eager VJP source has inconsistent inputs or outputs".into(),
3960 ));
3961 }
3962 let wrt_input_indices = wrts
3963 .iter()
3964 .map(|wrt| {
3965 let key = wrt.semantic_trace.as_ref()?.input_key()?;
3966 source.input_key_index(&key)
3967 })
3968 .collect::<Vec<_>>();
3969 let mut active_inputs = vec![false; source.input_count()];
3970 for &index in wrt_input_indices.iter().flatten() {
3971 active_inputs[index] = true;
3972 }
3973 if !active_inputs.iter().any(|&active| active) {
3974 return Ok(vec![None; wrts.len()]);
3975 }
3976
3977 // One transform and execution for the complete active set shares primal
3978 // work and cotangents between leaves instead of replaying them per target.
3979 // S2: check prepared-derivative cache before AD transform + compile_frozen.
3980 let target = ctx.ensure_engine(engine)?;
3981 let cache_key = PreparedDerivativeCacheKey {
3982 target: target.clone(),
3983 semantic_fingerprint: source.program().semantic_fingerprint(),
3984 runtime_epoch: ctx.runtime.epoch().map_err(|source| {
3985 Error::runtime_state_source("semantic_eager_vjp", ErrorPhase::Execution, source)
3986 })?,
3987 active_inputs: active_inputs.clone().into_boxed_slice(),
3988 input_metadata: source.frozen_program().input_metadata_with_bound_shapes(),
3989 };
3990 let prepared = { ctx.lock_prepared_derivative_cache()?.get(&cache_key) };
3991 let (
3992 seed_input_index,
3993 derivative_output_indices,
3994 derivative_program,
3995 execution_program,
3996 saved_input_indices,
3997 prepared_runtime,
3998 ) = if let Some(prepared) = prepared {
3999 (
4000 prepared.seed_input_index,
4001 prepared.derivative_output_indices.clone(),
4002 Arc::clone(&prepared.program),
4003 Arc::clone(&prepared.execution_program),
4004 prepared.saved_input_indices.clone(),
4005 Some(Arc::clone(&prepared.prepared)),
4006 )
4007 } else {
4008 let mut active_outputs = vec![false; source.output_count()];
4009 active_outputs[0] = true;
4010 let ad = AdContext::with_rules_and_transform_cache(
4011 ctx.semantic_extension_rules.clone(),
4012 Arc::clone(&ctx.ad_transform_cache),
4013 );
4014 let derivative = ad
4015 .vjp_program(source.frozen_program(), &active_inputs, &active_outputs)
4016 .map_err(|source| {
4017 Error::runtime_state_source("semantic_eager_vjp", ErrorPhase::GraphBuild, source)
4018 })?;
4019 let seed_input_index = derivative
4020 .derivative_input_indices()
4021 .first()
4022 .copied()
4023 .flatten();
4024 let Some(seed_input_index) = seed_input_index else {
4025 return Ok(vec![None; wrts.len()]);
4026 };
4027 let derivative_output_indices = derivative
4028 .derivative_output_indices()
4029 .to_vec()
4030 .into_boxed_slice();
4031 let program = Arc::new(compiler.compile_frozen_program(derivative.frozen())?);
4032 let (execution_program, saved_input_indices) = if saved.is_empty() {
4033 (Arc::clone(&program), Vec::new())
4034 } else {
4035 let (execution, indices) = crate::semantic_transform::semantic_vjp_with_saved_outputs(
4036 source.frozen_program(),
4037 &active_inputs,
4038 &active_outputs,
4039 &ctx.semantic_extension_rules,
4040 &(1..source.output_count()).collect::<Vec<_>>(),
4041 )
4042 .map_err(|source| {
4043 Error::runtime_state_source("semantic_eager_vjp", ErrorPhase::GraphBuild, source)
4044 })?;
4045 (
4046 Arc::new(compiler.compile_frozen_program(execution.frozen())?),
4047 indices,
4048 )
4049 };
4050 (
4051 seed_input_index,
4052 derivative_output_indices,
4053 program,
4054 execution_program,
4055 saved_input_indices,
4056 None,
4057 )
4058 };
4059
4060 let cotangent_tensor = Arc::new(RetainedValue::from_tensor(
4061 cotangent.to_tensor(&mut *engine)?,
4062 ));
4063 let input_count = execution_program.input_count();
4064 let mut owned_inputs: Vec<Option<Tensor>> = (0..input_count).map(|_| None).collect();
4065 // Residuals, primal bindings and the seed are staged on the caller's engine.
4066 (|| -> Result<()> {
4067 let session: &mut dyn Engine = &mut *engine;
4068 for (value, &index) in saved.iter().zip(&saved_input_indices) {
4069 let read = value.value.tensor_read("eager residual")?;
4070 let Some(slot) = owned_inputs.get_mut(index) else {
4071 return Err(Error::Internal(format!(
4072 "semantic eager VJP residual index {index} is outside {input_count} inputs"
4073 )));
4074 };
4075 *slot = Some(session.to_contiguous_read(read)?);
4076 }
4077 for (source_input_index, (_, tensor)) in source.bindings().iter().enumerate() {
4078 let Some(slot) = owned_inputs.get_mut(source_input_index) else {
4079 return Err(Error::Internal(format!(
4080 "semantic eager VJP derivative program has no primal input slot {source_input_index}"
4081 )));
4082 };
4083 *slot = Some(copy_value_in_session(session, tensor)?);
4084 }
4085 let Some(slot) = owned_inputs.get_mut(seed_input_index) else {
4086 return Err(Error::Internal(format!(
4087 "semantic eager VJP seed input index {seed_input_index} is outside {input_count} inputs"
4088 )));
4089 };
4090 *slot = Some(copy_value_in_session(session, cotangent_tensor.as_ref())?);
4091 Ok(())
4092 })()?;
4093 let input_refs = owned_inputs
4094 .iter()
4095 .enumerate()
4096 .map(|(index, tensor)| {
4097 tensor.as_ref().ok_or_else(|| {
4098 Error::Internal(format!(
4099 "semantic eager VJP derivative input {index} was not populated"
4100 ))
4101 })
4102 })
4103 .collect::<Result<Vec<_>>>()?;
4104 let input_reads = input_refs
4105 .iter()
4106 .map(|tensor| TensorRead::from_tensor(tensor))
4107 .collect::<Vec<_>>();
4108 let prepared_runtime = if let Some(prepared_runtime) = prepared_runtime {
4109 prepared_runtime
4110 } else {
4111 let prepared_runtime = Arc::new(ctx.runtime.prepare(
4112 &execution_program,
4113 &input_reads,
4114 std::slice::from_ref(&target),
4115 )?);
4116 let entry = Arc::new(PreparedDerivative {
4117 program: Arc::clone(&derivative_program),
4118 execution_program: Arc::clone(&execution_program),
4119 saved_input_indices: saved_input_indices.clone(),
4120 prepared: Arc::clone(&prepared_runtime),
4121 seed_input_index,
4122 derivative_output_indices: derivative_output_indices.clone(),
4123 });
4124 ctx.lock_prepared_derivative_cache()?
4125 .insert(cache_key, entry);
4126 prepared_runtime
4127 };
4128 let mut outputs = run_prepared_on(&ctx.runtime, engine, &prepared_runtime, &input_reads)?
4129 .into_iter()
4130 .map(Some)
4131 .collect::<Vec<_>>();
4132 let cotangent_trace =
4133 TracedTensor::from_shared_tensor_value_symbolic_shape(Arc::clone(&cotangent_tensor))?;
4134
4135 #[cfg(test)]
4136 EAGER_SEMANTIC_VJP_EXECUTIONS.fetch_add(1, Ordering::Relaxed);
4137
4138 wrts.iter()
4139 .zip(wrt_input_indices)
4140 .map(|(wrt, input_index)| {
4141 let Some(derivative_output_index) = input_index
4142 .and_then(|index| derivative_output_indices.get(index).copied().flatten())
4143 else {
4144 return Ok(None);
4145 };
4146 let result = outputs
4147 .get_mut(derivative_output_index)
4148 .and_then(Option::take)
4149 .ok_or_else(|| {
4150 Error::Internal(format!(
4151 "semantic eager VJP derivative output {derivative_output_index} unavailable"
4152 ))
4153 })?;
4154 let wrt_trace = wrt.semantic_trace.as_ref().ok_or_else(|| {
4155 Error::Internal("active eager VJP input has no semantic trace".into())
4156 })?;
4157 let semantic_trace = derivative_trace_from_frozen_program(
4158 &source,
4159 derivative_program.frozen_program(),
4160 derivative_output_index,
4161 &[(seed_input_index, Arc::clone(&cotangent_tensor))],
4162 &[&output_trace, wrt_trace, &cotangent_trace],
4163 None,
4164 "semantic_eager_vjp",
4165 )?;
4166 Ok(Some(EagerTensor::new_result_with_semantic_trace(
4167 Arc::clone(ctx),
4168 eager_val_key(),
4169 result,
4170 true,
4171 output.trace.clone(),
4172 Some(semantic_trace),
4173 )?))
4174 })
4175 .collect()
4176}
4177
4178fn semantic_eager_jvp_optional(
4179 ctx: &Arc<EagerRuntime>,
4180 engine: &mut dyn ExecutionEngine,
4181 output: &EagerTensor,
4182 wrt: &EagerTensor,
4183 tangent: &EagerTensor,
4184) -> Result<Option<EagerTensor>> {
4185 let (Some(raw_output_trace), Some(wrt_trace)) =
4186 (output.semantic_trace.as_ref(), wrt.semantic_trace.as_ref())
4187 else {
4188 return Ok(None);
4189 };
4190 let Some(wrt_key) = wrt_trace.input_key() else {
4191 return Ok(None);
4192 };
4193 if !raw_output_trace.has_attached_input_key(&wrt_key) {
4194 return Ok(None);
4195 }
4196
4197 // First AD request on this output: run the deferred graph analysis once
4198 // over the whole raw carrier chain before compiling.
4199 let output_trace = analyze_deferred_semantic_trace(raw_output_trace)?;
4200
4201 let mut compiler = GraphCompiler::new();
4202 let source = compile_ad_source(&mut compiler, &output_trace)?;
4203 if source.output_count() != 1
4204 || source.input_keys().len() != source.input_count()
4205 || source.bindings().len() != source.input_count()
4206 {
4207 return Err(Error::Internal(
4208 "semantic eager JVP source has inconsistent inputs or outputs".into(),
4209 ));
4210 }
4211 let Some(wrt_input_index) = source.input_key_index(&wrt_key) else {
4212 return Ok(None);
4213 };
4214
4215 let mut active_inputs = vec![false; source.input_count()];
4216 if let Some(active) = active_inputs.get_mut(wrt_input_index) {
4217 *active = true;
4218 } else {
4219 return Err(Error::Internal(
4220 "semantic eager JVP input index is outside the source inputs".into(),
4221 ));
4222 }
4223 let ad = AdContext::with_rules_and_transform_cache(
4224 ctx.semantic_extension_rules.clone(),
4225 Arc::clone(&ctx.ad_transform_cache),
4226 );
4227 let derivative = ad
4228 .jvp_program(source.frozen_program(), &active_inputs)
4229 .map_err(|source| {
4230 Error::runtime_state_source("semantic_eager_jvp", ErrorPhase::GraphBuild, source)
4231 })?;
4232 // derivative_input_indices maps source input → derivative seed input.
4233 let Some(seed_input_index) = derivative
4234 .derivative_input_indices()
4235 .get(wrt_input_index)
4236 .copied()
4237 .flatten()
4238 else {
4239 return Ok(None);
4240 };
4241 // derivative_output_indices maps source output → derivative output.
4242 // There is always exactly one source output (guarded above).
4243 let Some(derivative_output_index) = derivative
4244 .derivative_output_indices()
4245 .first()
4246 .copied()
4247 .flatten()
4248 else {
4249 return Ok(None);
4250 };
4251
4252 let derivative_program = compiler.compile_frozen_program(derivative.frozen())?;
4253 let tangent_tensor = Arc::new(RetainedValue::from_tensor(tangent.to_tensor(&mut *engine)?));
4254 let input_count = derivative_program.input_count();
4255 let mut owned_inputs: Vec<Option<Tensor>> = (0..input_count).map(|_| None).collect();
4256 // Primal bindings and the seed are staged on the caller's engine.
4257 (|| -> Result<()> {
4258 let session: &mut dyn Engine = &mut *engine;
4259 for (source_input_index, (_, tensor)) in source.bindings().iter().enumerate() {
4260 let Some(slot) = owned_inputs.get_mut(source_input_index) else {
4261 return Err(Error::Internal(format!(
4262 "semantic eager JVP derivative program has no primal input slot {source_input_index}"
4263 )));
4264 };
4265 *slot = Some(copy_value_in_session(session, tensor)?);
4266 }
4267 let Some(slot) = owned_inputs.get_mut(seed_input_index) else {
4268 return Err(Error::Internal(format!(
4269 "semantic eager JVP seed input index {seed_input_index} is outside {input_count} inputs"
4270 )));
4271 };
4272 *slot = Some(copy_value_in_session(session, tangent_tensor.as_ref())?);
4273 Ok(())
4274 })()?;
4275 let input_refs = owned_inputs
4276 .iter()
4277 .enumerate()
4278 .map(|(index, tensor)| {
4279 tensor.as_ref().ok_or_else(|| {
4280 Error::Internal(format!(
4281 "semantic eager JVP derivative input {index} was not populated"
4282 ))
4283 })
4284 })
4285 .collect::<Result<Vec<_>>>()?;
4286 let input_reads = input_refs
4287 .iter()
4288 .map(|tensor| TensorRead::from_tensor(tensor))
4289 .collect::<Vec<_>>();
4290 let outputs = ctx.run_program(engine, &derivative_program, &input_reads)?;
4291 let output_count = outputs.len();
4292 let Some(result) = outputs.into_iter().nth(derivative_output_index) else {
4293 return Err(Error::Internal(format!(
4294 "semantic eager JVP derivative output index {derivative_output_index} is outside {} outputs",
4295 output_count
4296 )));
4297 };
4298 let tangent_trace =
4299 TracedTensor::from_shared_tensor_value_symbolic_shape(Arc::clone(&tangent_tensor))?;
4300 let semantic_trace = derivative_trace_from_frozen_program(
4301 &source,
4302 derivative.frozen(),
4303 derivative_output_index,
4304 &[(seed_input_index, Arc::clone(&tangent_tensor))],
4305 &[&output_trace, wrt_trace, &tangent_trace],
4306 None,
4307 "semantic_eager_jvp",
4308 )?;
4309
4310 Ok(Some(EagerTensor::new_result_with_semantic_trace(
4311 Arc::clone(ctx),
4312 eager_val_key(),
4313 result,
4314 true,
4315 None,
4316 Some(semantic_trace),
4317 )?))
4318}
4319
4320fn validate_same_runtime(
4321 runtime: &Arc<EagerRuntime>,
4322 tensor: &EagerTensor,
4323 role: &'static str,
4324) -> Result<()> {
4325 if tensor.ctx_id() != runtime.id() {
4326 return Err(Error::ContextMismatch {
4327 lhs: runtime.id(),
4328 rhs: tensor.ctx_id(),
4329 });
4330 }
4331 let _ = role;
4332 Ok(())
4333}
4334
4335fn copy_value_in_session(session: &mut dyn Engine, value: &RetainedValue) -> Result<Tensor> {
4336 let read = value.tensor_read().map_err(|error| {
4337 Error::runtime_state_source("copy_value_for_runtime", ErrorPhase::Execution, error)
4338 })?;
4339 session.to_contiguous_read(read).map_err(Error::from)
4340}
4341
4342fn validate_seed_tensor(op: &'static str, primal: &EagerTensor, seed: &EagerTensor) -> Result<()> {
4343 if primal.dtype() != seed.dtype() {
4344 return Err(
4345 tenferro_tensor::Error::dtype_mismatch(op, primal.dtype(), seed.dtype()).into(),
4346 );
4347 }
4348 if primal.shape() != seed.shape() {
4349 return Err(
4350 tenferro_tensor::Error::shape_mismatch(op, primal.shape(), seed.shape()).into(),
4351 );
4352 }
4353 Ok(())
4354}
4355
4356/// Eager tensor with reverse-mode autodiff over concrete tensor values.
4357///
4358/// This executes each primitive immediately and records a lightweight reverse
4359/// DAG for `backward()`. Gradients accumulate across repeated `backward()`
4360/// calls until they are cleared explicitly.
4361///
4362/// # Examples
4363///
4364/// ```
4365/// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
4366///
4367/// let ctx = EagerRuntime::new()?;
4368/// let x = EagerTensor::requires_grad_in(Tensor::from_vec_col_major(vec![3], vec![1.0_f64, 2.0, 3.0]).unwrap(), ctx)?;
4369/// for _ in 0..2 {
4370/// let loss = x.runtime().with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
4371/// let squared = s.mul(&x, &x)?;
4372/// s.reduce_sum(&squared, Some(&[0]))
4373/// })?;
4374/// loss.backward(&mut tenferro_cpu::CpuEngine::serial())?;
4375/// }
4376///
4377/// assert_eq!(x.grad()?.unwrap().as_slice::<f64>().unwrap(), &[4.0, 8.0, 12.0]);
4378/// x.clear_grad()?;
4379///
4380/// assert!(x.grad().unwrap().is_none());
4381/// # Ok::<(), tenferro_ad::Error>(())
4382/// ```
4383#[derive(Clone)]
4384pub struct EagerTensor {
4385 pub(crate) key: ValueKey<StdTensorOp>,
4386 pub(crate) trace: Option<EagerTrace>,
4387 pub(crate) semantic_trace: Option<TracedTensor>,
4388 pub(crate) requires_grad: bool,
4389 grad_slot: GradSlot,
4390 pub(crate) ctx: Arc<EagerRuntime>,
4391 _record: Arc<EagerTensorRecord>,
4392}
4393
4394pub(crate) struct EagerTensorRecord {
4395 value: Arc<AdValueRecord>,
4396 key: ValueKey<StdTensorOp>,
4397 trace: Option<EagerTrace>,
4398 semantic_trace: Option<TracedTensor>,
4399 requires_grad: bool,
4400 grad_slot: GradSlot,
4401 ctx: Arc<EagerRuntime>,
4402}
4403
4404struct EagerTensorParts {
4405 ctx: Arc<EagerRuntime>,
4406 key: ValueKey<StdTensorOp>,
4407 requires_grad: bool,
4408 trace: Option<EagerTrace>,
4409 semantic_trace: Option<TracedTensor>,
4410 value: Arc<AdValueRecord>,
4411 register_value: bool,
4412}
4413
4414impl fmt::Debug for EagerTensor {
4415 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
4416 f.debug_struct("EagerTensor")
4417 .field("dtype", &self.dtype())
4418 .field("shape", &self.shape())
4419 .field("key", &self.key)
4420 .field("requires_grad", &self.requires_grad)
4421 .field("has_trace", &self.trace.is_some())
4422 .field("has_semantic_trace", &self.semantic_trace.is_some())
4423 .field("ctx_id", &self.ctx_id())
4424 .finish_non_exhaustive()
4425 }
4426}
4427
4428impl EagerTensor {
4429 /// Create an untracked eager tensor inside an existing eager context.
4430 ///
4431 /// # Examples
4432 ///
4433 /// ```
4434 /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
4435 ///
4436 /// let ctx = EagerRuntime::new()?;
4437 /// let x = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0]).unwrap(), ctx)?;
4438 ///
4439 /// assert_eq!(x.value()?.as_slice::<f64>().unwrap(), &[1.0, 2.0]);
4440 /// # Ok::<(), tenferro_ad::Error>(())
4441 /// ```
4442 ///
4443 /// # Errors
4444 ///
4445 /// Returns [`tenferro_runtime::Error::RuntimeState`] when metadata cannot
4446 /// be registered in the target context, or a typed tensor/backend error
4447 /// while materializing the source value.
4448 pub fn from_tensor_in(tensor: Tensor, ctx: Arc<EagerRuntime>) -> Result<Self> {
4449 Self::new_leaf(ctx, tensor, false)
4450 }
4451
4452 /// Create an untracked eager tensor from compact column-major data inside
4453 /// an existing eager runtime.
4454 ///
4455 /// # Errors
4456 ///
4457 /// Returns [`Error::TensorRuntime`] with
4458 /// [`tenferro_tensor::ValidationError::ShapeMismatch`] when the shape and
4459 /// data length disagree, or with
4460 /// [`tenferro_tensor::ValidationError::IntegerOverflow`] when shape
4461 /// arithmetic overflows. Returns [`Error::RuntimeState`] when eager
4462 /// metadata cannot be registered.
4463 pub fn from_vec_col_major_in<T: TensorScalar>(
4464 shape: impl IntoShapeVec,
4465 data: Vec<T>,
4466 ctx: Arc<EagerRuntime>,
4467 ) -> Result<Self> {
4468 Self::from_tensor_in(Tensor::from_vec_col_major(shape, data)?, ctx)
4469 }
4470
4471 /// Create a tracked eager leaf inside an existing eager context.
4472 ///
4473 /// # Examples
4474 ///
4475 /// ```
4476 /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
4477 ///
4478 /// let ctx = EagerRuntime::new()?;
4479 /// let x = EagerTensor::requires_grad_in(Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0]).unwrap(), ctx)?;
4480 ///
4481 /// assert!(x.grad().unwrap().is_none());
4482 /// # Ok::<(), tenferro_ad::Error>(())
4483 /// ```
4484 ///
4485 /// # Errors
4486 ///
4487 /// Returns [`tenferro_runtime::Error::RuntimeState`] when gradient metadata
4488 /// cannot be registered in the target context, or a typed tensor/backend
4489 /// error while creating the leaf.
4490 pub fn requires_grad_in(tensor: Tensor, ctx: Arc<EagerRuntime>) -> Result<Self> {
4491 Self::new_leaf(ctx, tensor, true)
4492 }
4493
4494 pub(crate) fn new_leaf(
4495 ctx: Arc<EagerRuntime>,
4496 tensor: Tensor,
4497 requires_grad: bool,
4498 ) -> Result<Self> {
4499 Self::new_leaf_with_session(ctx, tensor, requires_grad, None)
4500 }
4501
4502 pub(crate) fn new_leaf_in_session(
4503 ctx: Arc<EagerRuntime>,
4504 tensor: Tensor,
4505 requires_grad: bool,
4506 session: &mut dyn Engine,
4507 ) -> Result<Self> {
4508 Self::new_leaf_with_session(ctx, tensor, requires_grad, Some(session))
4509 }
4510
4511 fn new_leaf_with_session(
4512 ctx: Arc<EagerRuntime>,
4513 tensor: Tensor,
4514 requires_grad: bool,
4515 session: Option<&mut dyn Engine>,
4516 ) -> Result<Self> {
4517 let key = eager_val_key();
4518 // A host-placement tensor needs no backend session: the CPU backend only
4519 // copies the host buffer for it, so entering a session would add
4520 // admission, provider exclusion, and session construction without doing
4521 // any provider work (#1704). Views and backend-family reads keep the
4522 // session path, and device runtimes keep it for every read.
4523 let read = TensorRead::from_tensor(&tensor);
4524 let semantic_tensor =
4525 match session {
4526 Some(session) => session.to_contiguous_read(read).map_err(Error::from)?,
4527 None => match cpu_host_owned_read(&read) {
4528 Some(materialized) => materialized.map_err(Error::from)?,
4529 None => return Err(Error::unsupported(
4530 "EagerTensor::new_leaf",
4531 ErrorPhase::Execution,
4532 "this tensor needs an engine to become an eager leaf; create it through \
4533 EagerSession::constant_from or EagerSession::variable_from",
4534 )),
4535 },
4536 };
4537 let semantic_value = Arc::new(RetainedValue::from_tensor(semantic_tensor));
4538 let semantic_trace = TracedTensor::from_shared_tensor_value_symbolic_shape(semantic_value)?;
4539 // Deferred materialization: the per-op/leaf global-metadata registry
4540 // write for the eager tensor key was unreadable after the semantic
4541 // trace became the sole AD carrier, so it is dropped.
4542 // ponytail: leaf input-key metadata is still registered by
4543 // `from_shared_tensor_value_symbolic_shape`; the eager-key entry was
4544 // vestigial and removed. Add back only if something reads it.
4545 let value = AdValueRecord::from_tensor(tensor, "EagerTensor::new_leaf")?;
4546 Self::from_parts(EagerTensorParts {
4547 ctx,
4548 key,
4549 requires_grad,
4550 trace: None,
4551 semantic_trace: Some(semantic_trace),
4552 value,
4553 register_value: true,
4554 })
4555 }
4556
4557 pub(crate) fn new_result(
4558 ctx: Arc<EagerRuntime>,
4559 key: ValueKey<StdTensorOp>,
4560 tensor: Tensor,
4561 requires_grad: bool,
4562 trace: Option<EagerTrace>,
4563 ) -> Result<Self> {
4564 Self::new_result_with_semantic_trace(ctx, key, tensor, requires_grad, trace, None)
4565 }
4566
4567 pub(crate) fn new_result_with_semantic_trace(
4568 ctx: Arc<EagerRuntime>,
4569 key: ValueKey<StdTensorOp>,
4570 tensor: Tensor,
4571 requires_grad: bool,
4572 trace: Option<EagerTrace>,
4573 semantic_trace: Option<TracedTensor>,
4574 ) -> Result<Self> {
4575 let value = AdValueRecord::from_tensor(tensor, "EagerTensor::new_result")?;
4576 Self::from_parts(EagerTensorParts {
4577 ctx,
4578 key,
4579 requires_grad,
4580 trace,
4581 semantic_trace,
4582 value,
4583 register_value: true,
4584 })
4585 }
4586
4587 pub(crate) fn new_unregistered_result_with_semantic_trace(
4588 ctx: Arc<EagerRuntime>,
4589 key: ValueKey<StdTensorOp>,
4590 tensor: Tensor,
4591 requires_grad: bool,
4592 trace: Option<EagerTrace>,
4593 semantic_trace: Option<TracedTensor>,
4594 ) -> Result<Self> {
4595 let value = AdValueRecord::from_tensor(tensor, "EagerTensor::new_unregistered_result")?;
4596 Self::from_parts(EagerTensorParts {
4597 ctx,
4598 key,
4599 requires_grad,
4600 trace,
4601 semantic_trace,
4602 value,
4603 register_value: false,
4604 })
4605 }
4606
4607 pub(crate) fn new_result_value(
4608 ctx: Arc<EagerRuntime>,
4609 key: ValueKey<StdTensorOp>,
4610 value: TensorValue,
4611 requires_grad: bool,
4612 trace: Option<EagerTrace>,
4613 semantic_trace: Option<TracedTensor>,
4614 ) -> Result<Self> {
4615 let (group, slot, dtype, shape) = value.try_into_group_parts().map_err(|_| {
4616 Error::runtime_state(
4617 "EagerTensor::new_result_value",
4618 ErrorPhase::Execution,
4619 "a TensorValue could not be transferred into its allocation group",
4620 )
4621 })?;
4622 let value = AdValueRecord::from_group(group, slot, dtype, shape);
4623 Self::from_parts(EagerTensorParts {
4624 ctx,
4625 key,
4626 requires_grad,
4627 trace,
4628 semantic_trace,
4629 value,
4630 register_value: true,
4631 })
4632 }
4633
4634 fn from_parts(parts: EagerTensorParts) -> Result<Self> {
4635 let EagerTensorParts {
4636 ctx,
4637 key,
4638 requires_grad,
4639 trace,
4640 semantic_trace,
4641 value,
4642 register_value,
4643 } = parts;
4644 let grad_slot = Arc::new(Mutex::new(None));
4645 if requires_grad {
4646 ctx.try_register_grad_slot(&key, &grad_slot)?;
4647 }
4648 let record = Arc::new(EagerTensorRecord {
4649 value: Arc::clone(&value),
4650 key: key.clone(),
4651 trace: trace.clone(),
4652 semantic_trace: semantic_trace.clone(),
4653 requires_grad,
4654 grad_slot: Arc::clone(&grad_slot),
4655 ctx: Arc::clone(&ctx),
4656 });
4657 if register_value {
4658 ctx.try_register_value_record(&key, &record)?;
4659 }
4660
4661 Ok(Self {
4662 key,
4663 trace,
4664 semantic_trace,
4665 requires_grad,
4666 grad_slot,
4667 ctx,
4668 _record: record,
4669 })
4670 }
4671
4672 pub(crate) fn new_untracked_result(ctx: Arc<EagerRuntime>, tensor: Tensor) -> Result<Self> {
4673 let value =
4674 AdValueRecord::from_untracked_tensor(tensor, "EagerTensor::new_untracked_result")?;
4675 Ok(Self::new_untracked_value_record(ctx, value, None))
4676 }
4677
4678 pub(crate) fn new_untracked_value_result(
4679 ctx: Arc<EagerRuntime>,
4680 value: TensorValue,
4681 ) -> Result<Self> {
4682 Self::new_untracked_value_result_with_semantic_trace(ctx, value, None)
4683 }
4684
4685 pub(crate) fn new_untracked_value_result_with_semantic_trace(
4686 ctx: Arc<EagerRuntime>,
4687 value: TensorValue,
4688 semantic_trace: Option<TracedTensor>,
4689 ) -> Result<Self> {
4690 let (group, slot, dtype, shape) = value.try_into_group_parts().map_err(|_| {
4691 Error::runtime_state(
4692 "EagerTensor::new_untracked_value_result",
4693 ErrorPhase::Execution,
4694 "a TensorValue could not be transferred into its allocation group",
4695 )
4696 })?;
4697 let value = AdValueRecord::from_group(group, slot, dtype, shape);
4698 Ok(Self::new_untracked_value_record(ctx, value, semantic_trace))
4699 }
4700
4701 fn new_untracked_value_record(
4702 ctx: Arc<EagerRuntime>,
4703 value: Arc<AdValueRecord>,
4704 semantic_trace: Option<TracedTensor>,
4705 ) -> Self {
4706 let key = eager_val_key();
4707 let grad_slot = Arc::new(Mutex::new(None));
4708 let record = Arc::new(EagerTensorRecord {
4709 value,
4710 key: key.clone(),
4711 trace: None,
4712 semantic_trace: semantic_trace.clone(),
4713 requires_grad: false,
4714 grad_slot: Arc::clone(&grad_slot),
4715 ctx: Arc::clone(&ctx),
4716 });
4717 Self {
4718 key,
4719 trace: None,
4720 semantic_trace,
4721 requires_grad: false,
4722 grad_slot,
4723 ctx,
4724 _record: record,
4725 }
4726 }
4727
4728 pub(crate) fn from_record(record: Arc<EagerTensorRecord>) -> Self {
4729 Self {
4730 key: record.key.clone(),
4731 trace: record.trace.clone(),
4732 semantic_trace: record.semantic_trace.clone(),
4733 requires_grad: record.requires_grad,
4734 grad_slot: Arc::clone(&record.grad_slot),
4735 ctx: Arc::clone(&record.ctx),
4736 _record: record,
4737 }
4738 }
4739
4740 /// Detach this tensor from the reverse graph.
4741 ///
4742 /// The returned tensor keeps the concrete value but no longer contributes
4743 /// gradients to the original graph.
4744 ///
4745 /// # Errors
4746 ///
4747 /// Returns [`Error::TensorRuntime`] when `engine` cannot materialize the
4748 /// retained value, [`Error::RuntimeStateSource`] when the value has no
4749 /// readable tensor or symbolic graph metadata registration is unavailable,
4750 /// and [`Error::RuntimeState`] when the detached record's allocation-group
4751 /// descriptor is invalid.
4752 ///
4753 /// # Examples
4754 ///
4755 /// ```
4756 /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
4757 ///
4758 /// let ctx = EagerRuntime::new()?;
4759 /// let x = EagerTensor::requires_grad_in(Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0]).unwrap(), ctx)?;
4760 /// let y = x.detach(&mut tenferro_cpu::CpuEngine::serial()).unwrap();
4761 ///
4762 /// assert_eq!(y.value()?.as_slice::<f64>().unwrap(), &[1.0, 2.0]);
4763 /// assert!(y.grad().unwrap().is_none());
4764 /// # Ok::<(), tenferro_ad::Error>(())
4765 /// ```
4766 pub fn detach(&self, engine: &mut dyn Engine) -> Result<Self> {
4767 let tensor = self.duplicate_value(engine)?;
4768 let semantic_trace = TracedTensor::from_tensor_symbolic_shape(tensor)?;
4769 Ok(Self::new_untracked_value_record(
4770 self.ctx.clone(),
4771 Arc::clone(&self._record.value),
4772 Some(semantic_trace),
4773 ))
4774 }
4775
4776 /// Detach this tensor from its graph and re-register it in a different
4777 /// context as an untracked leaf.
4778 ///
4779 /// # Examples
4780 ///
4781 /// ```
4782 /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
4783 ///
4784 /// let ctx_a = EagerRuntime::new()?;
4785 /// let ctx_b = EagerRuntime::new()?;
4786 /// let x = EagerTensor::requires_grad_in(Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0]).unwrap(), ctx_a)?;
4787 /// let d = x.detach_into(&mut tenferro_cpu::CpuEngine::serial(), &ctx_b)?;
4788 ///
4789 /// assert!(!d.tracks_grad());
4790 /// assert_eq!(d.ctx_id(), ctx_b.id());
4791 /// # Ok::<(), tenferro_ad::Error>(())
4792 /// ```
4793 ///
4794 /// # Errors
4795 ///
4796 /// Returns [`Error::RuntimeState`] if the source cannot be materialized or
4797 /// the target context cannot register its metadata.
4798 pub fn detach_into(&self, engine: &mut dyn Engine, ctx: &Arc<EagerRuntime>) -> Result<Self> {
4799 let tensor = self.to_tensor(&mut *engine)?;
4800 Self::new_leaf_in_session(Arc::clone(ctx), tensor, false, engine)
4801 }
4802
4803 /// Borrow the retained value without creating an owner or copy.
4804 ///
4805 /// # Errors
4806 ///
4807 /// Returns [`Error::RuntimeState`] when the retained allocation-group
4808 /// descriptor is unavailable or invalid.
4809 pub fn value(&self) -> Result<ValueGuard<'_>> {
4810 self._record.value.value("EagerTensor::value")
4811 }
4812
4813 /// Explicitly duplicate this value into a fresh standalone allocation.
4814 ///
4815 /// # Errors
4816 ///
4817 /// Returns [`Error::RuntimeState`] when the retained value or execution
4818 /// session is unavailable, or a typed host/backend error when the value
4819 /// cannot be materialized as a contiguous tensor.
4820 ///
4821 /// # Examples
4822 ///
4823 /// ```
4824 /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
4825 ///
4826 /// let ctx = EagerRuntime::new()?;
4827 /// let value = EagerTensor::from_tensor_in(
4828 /// Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0])?,
4829 /// ctx,
4830 /// )?;
4831 /// let duplicate = value.duplicate_value(&mut tenferro_cpu::CpuEngine::serial())?;
4832 /// assert_eq!(duplicate.as_slice::<f64>()?, &[1.0, 2.0]);
4833 /// # Ok::<(), tenferro_ad::Error>(())
4834 /// ```
4835 pub fn duplicate_value(&self, engine: &mut dyn Engine) -> Result<Tensor> {
4836 self.duplicate_value_in_session(engine)
4837 }
4838
4839 fn duplicate_host_value(&self) -> Option<Tensor> {
4840 // A pooled value duplicates through its descriptor view; a caller-owned
4841 // payload has no such view and duplicates through its own read path, which
4842 // copies the value while keeping its element type.
4843 self.value().ok()?.duplicate_host_tensor().ok()
4844 }
4845
4846 pub(crate) fn duplicate_value_in_session(&self, session: &mut dyn Engine) -> Result<Tensor> {
4847 if let Some(tensor) = self.duplicate_host_value() {
4848 return Ok(tensor);
4849 }
4850 let read = self
4851 ._record
4852 .value
4853 .tensor_read("EagerTensor::duplicate_value")?;
4854 session.to_contiguous_read(read).map_err(Error::from)
4855 }
4856
4857 // INVARIANT: the error variants return the unchanged eager handle so a
4858 // caller can retry ownership extraction without an implicit copy.
4859 #[allow(clippy::result_large_err)]
4860 /// Consume this handle and structurally extract its retained allocation.
4861 ///
4862 /// A shared handle is returned unchanged as [`IntoValueError::NotUnique`].
4863 /// Group extraction failures return the unchanged handle and typed group
4864 /// error; no copy or fallback materialization is attempted.
4865 ///
4866 /// # Errors
4867 ///
4868 /// Returns [`IntoValueError::NotUnique`] when another handle retains the
4869 /// value, or [`IntoValueError::Extract`] when structural group extraction
4870 /// fails because the allocation is aliased or its descriptor is invalid.
4871 ///
4872 /// # Examples
4873 ///
4874 /// ```
4875 /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
4876 ///
4877 /// let ctx = EagerRuntime::new()?;
4878 /// let value = EagerTensor::from_tensor_in(
4879 /// Tensor::from_vec_col_major(vec![1], vec![3.0_f64])?,
4880 /// ctx,
4881 /// )?;
4882 /// let owner = value
4883 /// .into_value()
4884 /// .expect("a uniquely owned value should be extractable");
4885 /// assert_eq!(owner.as_slice::<f64>()?, &[3.0]);
4886 /// # Ok::<(), tenferro_ad::Error>(())
4887 /// ```
4888 pub fn into_value(self) -> std::result::Result<Tensor, IntoValueError<Self>> {
4889 if Arc::strong_count(&self._record) != 1 {
4890 return Err(IntoValueError::NotUnique(self));
4891 }
4892 let Self { _record, .. } = self;
4893 let record = match Arc::try_unwrap(_record) {
4894 Ok(record) => record,
4895 Err(record) => return Err(IntoValueError::NotUnique(Self::from_record(record))),
4896 };
4897 let EagerTensorRecord {
4898 value,
4899 key,
4900 trace,
4901 semantic_trace,
4902 requires_grad,
4903 grad_slot,
4904 ctx,
4905 } = record;
4906 let value = match Arc::try_unwrap(value) {
4907 Ok(value) => value,
4908 Err(value) => {
4909 let record = Arc::new(EagerTensorRecord {
4910 value,
4911 key,
4912 trace,
4913 semantic_trace,
4914 requires_grad,
4915 grad_slot,
4916 ctx,
4917 });
4918 return Err(IntoValueError::NotUnique(Self::from_record(record)));
4919 }
4920 };
4921 let AdValueRecord {
4922 container,
4923 dtype,
4924 shape,
4925 } = value;
4926 let container = match Arc::try_unwrap(container) {
4927 Ok(container) => container,
4928 Err(container) => {
4929 let record = Arc::new(EagerTensorRecord {
4930 value: Arc::new(AdValueRecord {
4931 container,
4932 dtype,
4933 shape,
4934 }),
4935 key,
4936 trace,
4937 semantic_trace,
4938 requires_grad,
4939 grad_slot,
4940 ctx,
4941 });
4942 return Err(IntoValueError::NotUnique(Self::from_record(record)));
4943 }
4944 };
4945 let (group, slot) = match container {
4946 // A caller-owned payload is handed back to its owner unchanged.
4947 RetentionContainer::CallerOwned { tensor } => return Ok(*tensor),
4948 RetentionContainer::Owned { tensor } => return Ok(tensor),
4949 RetentionContainer::Pooled { group, slot } => (group, slot),
4950 };
4951 match group.into_tensor(slot) {
4952 Ok(tensor) => Ok(tensor),
4953 Err((group, error)) => {
4954 let record = Arc::new(EagerTensorRecord {
4955 value: Arc::new(AdValueRecord {
4956 container: Arc::new(RetentionContainer::Pooled { group, slot }),
4957 dtype,
4958 shape,
4959 }),
4960 key,
4961 trace,
4962 semantic_trace,
4963 requires_grad,
4964 grad_slot,
4965 ctx,
4966 });
4967 Err(IntoValueError::Extract {
4968 value: Self::from_record(record),
4969 error,
4970 })
4971 }
4972 }
4973 }
4974
4975 /// Return this tensor's scalar dtype without materializing through
4976 /// [`value`](Self::value).
4977 pub fn dtype(&self) -> DType {
4978 self._record.value.dtype()
4979 }
4980
4981 /// Return this tensor's logical shape without materializing through
4982 /// [`value`](Self::value).
4983 pub fn shape(&self) -> &[usize] {
4984 self._record.value.shape()
4985 }
4986
4987 /// Borrow this tensor value as a [`TensorRead`].
4988 ///
4989 /// This is the preferred borrowed input boundary for executor calls. It
4990 /// preserves the option to replace eager storage with non-contiguous views
4991 /// without forcing callers through [`value`](Self::value).
4992 ///
4993 /// # Panics
4994 ///
4995 /// Panics if a validated eager value record becomes unavailable, which
4996 /// indicates an internal invariant violation.
4997 pub fn tensor_read(&self) -> TensorRead<'_> {
4998 self._record
4999 .value
5000 .tensor_read("EagerTensor::tensor_read")
5001 .expect("validated eager value record")
5002 }
5003
5004 /// Materialize this eager tensor as an owned [`Tensor`].
5005 ///
5006 /// This is the owned materialization boundary for callers that need a
5007 /// standalone compact tensor. The operation is fallible because eager
5008 /// values may be backed by lazy or backend-resident storage.
5009 ///
5010 /// # Errors
5011 ///
5012 /// Returns [`Error::RuntimeState`] if backend state is unavailable, or a
5013 /// typed tensor backend error when contiguous materialization fails.
5014 pub fn to_tensor(&self, engine: &mut dyn Engine) -> Result<Tensor> {
5015 self.duplicate_value(engine)
5016 }
5017
5018 /// Return the accumulated gradient currently stored for this tensor.
5019 ///
5020 /// The stored gradient accumulates across repeated `backward()` calls
5021 /// until it is cleared explicitly.
5022 ///
5023 /// For complex scalar losses, stored gradients use tenferro's
5024 /// Hermitian-adjoint cotangent convention. See
5025 /// <https://tensor4all.org/tenferro-rs/guides/complex-ad.html>.
5026 ///
5027 /// # Examples
5028 ///
5029 /// ```
5030 /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
5031 ///
5032 /// let ctx = EagerRuntime::new()?;
5033 /// let x = EagerTensor::requires_grad_in(Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0]).unwrap(), ctx.clone()).unwrap();
5034 /// let loss = ctx.with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
5035 /// let y = s.exp(&x)?;
5036 /// s.reduce_sum(&y, Some(&[0]))
5037 /// })?;
5038 /// let _cotangents = loss.backward(&mut tenferro_cpu::CpuEngine::serial()).unwrap();
5039 ///
5040 /// let grad = x.grad()?.unwrap();
5041 /// assert_eq!(grad.shape(), &[2]);
5042 /// # Ok::<(), tenferro_ad::Error>(())
5043 /// ```
5044 ///
5045 /// # Errors
5046 ///
5047 /// Returns [`Error::RuntimeState`] if the gradient slot is poisoned or no
5048 /// longer available.
5049 pub fn grad(&self) -> Result<Option<GradientValue>> {
5050 self.grad_slot
5051 .lock()
5052 .map_err(|_| {
5053 Error::runtime_state(
5054 "eager_gradient_slot",
5055 ErrorPhase::Execution,
5056 "lock poisoned",
5057 )
5058 })
5059 .map(|slot| {
5060 slot.as_ref().map(|record| GradientValue {
5061 record: Arc::clone(record),
5062 })
5063 })
5064 }
5065
5066 /// Clear the accumulated gradient stored for this tensor.
5067 ///
5068 /// This only affects this tensor's gradient slot. Other tensors in the
5069 /// same context retain their gradients until they are cleared explicitly or
5070 /// overwritten by later accumulation.
5071 ///
5072 /// # Examples
5073 ///
5074 /// ```
5075 /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
5076 ///
5077 /// let ctx = EagerRuntime::new()?;
5078 /// let x = EagerTensor::requires_grad_in(Tensor::from_vec_col_major(vec![3], vec![1.0_f64, 2.0, 3.0]).unwrap(), ctx.clone()).unwrap();
5079 /// let y = EagerTensor::requires_grad_in(Tensor::from_vec_col_major(vec![3], vec![4.0_f64, 5.0, 6.0]).unwrap(), ctx).unwrap();
5080 /// let loss = x.runtime().with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
5081 /// let product = s.mul(&x, &y)?;
5082 /// s.reduce_sum(&product, Some(&[0]))
5083 /// })?;
5084 /// let _ = loss.backward(&mut tenferro_cpu::CpuEngine::serial()).unwrap();
5085 ///
5086 /// x.clear_grad()?;
5087 ///
5088 /// assert!(x.grad()?.is_none());
5089 /// assert!(y.grad()?.is_some());
5090 /// # Ok::<(), tenferro_ad::Error>(())
5091 /// ```
5092 ///
5093 /// # Errors
5094 ///
5095 /// Returns [`Error::RuntimeState`] if the gradient slot lock is poisoned.
5096 pub fn clear_grad(&self) -> Result<()> {
5097 *self.grad_slot.lock().map_err(|_| {
5098 Error::runtime_state(
5099 "eager_gradient_slot",
5100 ErrorPhase::Execution,
5101 "lock poisoned",
5102 )
5103 })? = None;
5104 Ok(())
5105 }
5106
5107 /// Report whether this tensor participates in gradient tracking.
5108 ///
5109 /// Tracked tensors keep a gradient slot in their eager context; untracked
5110 /// tensors and detached tensors do not.
5111 ///
5112 /// # Examples
5113 ///
5114 /// ```
5115 /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
5116 ///
5117 /// let ctx = EagerRuntime::new()?;
5118 /// let plain = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0]).unwrap(), ctx.clone()).unwrap();
5119 /// let tracked = EagerTensor::requires_grad_in(Tensor::from_vec_col_major(vec![2], vec![3.0_f64, 4.0]).unwrap(), ctx.clone()).unwrap();
5120 /// let detached = tracked.detach(&mut tenferro_cpu::CpuEngine::serial()).unwrap();
5121 ///
5122 /// assert!(!plain.tracks_grad());
5123 /// assert!(tracked.tracks_grad());
5124 /// assert!(!detached.tracks_grad());
5125 /// # Ok::<(), tenferro_ad::Error>(())
5126 /// ```
5127 pub fn tracks_grad(&self) -> bool {
5128 self.requires_grad
5129 }
5130
5131 #[cfg(test)]
5132 fn debug_trace_saved_value_count(&self) -> Option<usize> {
5133 None
5134 }
5135
5136 /// Return the opaque identifier of the context this tensor belongs to.
5137 ///
5138 /// # Examples
5139 ///
5140 /// ```
5141 /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
5142 ///
5143 /// let ctx = EagerRuntime::new()?;
5144 /// let x = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![1], vec![1.0_f64]).unwrap(), ctx.clone()).unwrap();
5145 ///
5146 /// assert_eq!(x.ctx_id(), ctx.id());
5147 /// # Ok::<(), tenferro_ad::Error>(())
5148 /// ```
5149 pub fn ctx_id(&self) -> ContextId {
5150 self.ctx.id()
5151 }
5152
5153 /// Borrow the eager runtime context that owns this tensor.
5154 pub fn runtime(&self) -> &Arc<EagerRuntime> {
5155 &self.ctx
5156 }
5157
5158 /// Check whether two tensors belong to the same eager context.
5159 ///
5160 /// # Examples
5161 ///
5162 /// ```
5163 /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
5164 ///
5165 /// let ctx = EagerRuntime::new()?;
5166 /// let x = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![1], vec![1.0_f64]).unwrap(), ctx.clone()).unwrap();
5167 /// let y = EagerTensor::from_tensor_in(Tensor::from_vec_col_major(vec![1], vec![2.0_f64]).unwrap(), ctx).unwrap();
5168 ///
5169 /// assert!(x.same_context(&y));
5170 /// # Ok::<(), tenferro_ad::Error>(())
5171 /// ```
5172 pub fn same_context(&self, other: &Self) -> bool {
5173 self.ctx_id() == other.ctx_id()
5174 }
5175
5176 #[cfg(test)]
5177 pub(crate) fn standard_graph_op(
5178 engine: &mut dyn ExecutionEngine,
5179 inputs: &[&Self],
5180 build_graph: impl FnOnce(&[TensorInputKey]) -> Result<Arc<Graph<StdTensorOp>>>,
5181 ) -> Result<Vec<Self>> {
5182 let Some(first) = inputs.first() else {
5183 return Err(Error::Internal(
5184 "standard eager graph op requires at least one input tensor".to_string(),
5185 ));
5186 };
5187 let ctx = Arc::clone(&first.ctx);
5188 for tensor in inputs.iter().skip(1) {
5189 if !first.same_context(tensor) {
5190 return Err(Error::ContextMismatch {
5191 lhs: first.ctx_id(),
5192 rhs: tensor.ctx_id(),
5193 });
5194 }
5195 }
5196
5197 let graph_input_keys = (0..inputs.len())
5198 .map(|_| next_input_key())
5199 .collect::<Vec<_>>();
5200 let graph = build_graph(&graph_input_keys)?;
5201 let initial_data = graph_input_keys
5202 .iter()
5203 .zip(inputs.iter())
5204 .map(|(key, tensor)| {
5205 Ok((
5206 ValueKey::Input(key.clone()),
5207 tensor.to_tensor(&mut *engine)?,
5208 ))
5209 })
5210 .collect::<Result<HashMap<_, _>>>()?;
5211 let execution = ctx.exec_standard_graph_outputs(engine, graph.as_ref(), initial_data)?;
5212 if execution.outputs.len() != graph.outputs().len() {
5213 return Err(Error::Internal(format!(
5214 "standard eager graph op expected {} graph outputs, got {}",
5215 graph.outputs().len(),
5216 execution.outputs.len()
5217 )));
5218 }
5219
5220 if !eager_grad_recording_enabled() || !inputs.iter().any(|input| input.requires_grad) {
5221 return execution
5222 .outputs
5223 .into_iter()
5224 .map(|output| {
5225 Self::new_unregistered_result_with_semantic_trace(
5226 Arc::clone(&ctx),
5227 eager_val_key(),
5228 output,
5229 false,
5230 None,
5231 None,
5232 )
5233 })
5234 .collect();
5235 }
5236
5237 let recorded = record_eager_graph_outputs(
5238 graph.as_ref(),
5239 &graph_input_keys,
5240 &execution.outputs,
5241 inputs,
5242 )?;
5243 if recorded.traces.len() != execution.outputs.len() {
5244 return Err(Error::Internal(format!(
5245 "standard eager graph op expected {} eager traces, got {}",
5246 execution.outputs.len(),
5247 recorded.traces.len()
5248 )));
5249 }
5250
5251 recorded
5252 .traces
5253 .into_iter()
5254 .zip(recorded.semantic_traces)
5255 .zip(execution.outputs)
5256 .map(|((trace, semantic_trace), output)| {
5257 Self::new_result_with_semantic_trace(
5258 Arc::clone(&ctx),
5259 trace.key,
5260 output,
5261 trace.requires_grad,
5262 trace.trace,
5263 semantic_trace,
5264 )
5265 })
5266 .collect()
5267 }
5268
5269 /// Run reverse-mode AD from this scalar output.
5270 ///
5271 /// Returns the full cotangent map produced by the reverse pass and also
5272 /// accumulates into `grad()` for tracked eager leaves reachable from this
5273 /// output. A pass with no gradients is an error, not an empty successful
5274 /// bundle. If a reachable leaf has no cotangent flow, the pass fails before
5275 /// accumulating any gradients; use [`EagerRuntime::vjp_optional`] to observe
5276 /// absent flow explicitly.
5277 ///
5278 /// For complex scalar outputs, cotangents use tenferro's Hermitian
5279 /// real-inner-product convention. See
5280 /// <https://tensor4all.org/tenferro-rs/guides/complex-ad.html>.
5281 ///
5282 /// # Examples
5283 ///
5284 /// ```
5285 /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
5286 ///
5287 /// let ctx = EagerRuntime::new()?;
5288 /// let x = EagerTensor::requires_grad_in(Tensor::from_vec_col_major(vec![3], vec![1.0_f64, 2.0, 3.0]).unwrap(), ctx).unwrap();
5289 /// for _ in 0..2 {
5290 /// let loss = x.runtime().with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| {
5291 /// let doubled = s.add(&x, &x)?;
5292 /// s.reduce_sum(&doubled, Some(&[0]))
5293 /// })?;
5294 /// loss.backward(&mut tenferro_cpu::CpuEngine::serial())?;
5295 /// }
5296 ///
5297 /// assert_eq!(x.grad().unwrap().unwrap().as_slice::<f64>().unwrap(), &[4.0, 4.0, 4.0]);
5298 /// # Ok::<(), tenferro_ad::Error>(())
5299 /// ```
5300 ///
5301 /// # Errors
5302 ///
5303 /// Returns [`Error::NonScalarGrad`] when this output is not scalar,
5304 /// [`Error::Validation`] when a reachable tracked leaf has no cotangent,
5305 /// [`Error::RuntimeState`] when no tracked leaf produces a gradient,
5306 /// [`Error::UnsupportedAdRule`] when a graph operation lacks a reverse rule,
5307 /// or a typed validation/backend/runtime-state error during the reverse pass.
5308 pub fn backward(&self, engine: &mut dyn ExecutionEngine) -> Result<Gradients> {
5309 if !self.shape().is_empty() {
5310 return Err(Error::NonScalarGrad {
5311 shape: self.shape().to_vec(),
5312 });
5313 }
5314
5315 let value = self.to_tensor(&mut *engine)?;
5316 let seed = one_like_tensor(&value, &mut *engine)?;
5317 self.backward_from_seed(engine, seed)
5318 }
5319
5320 /// Run reverse-mode AD from this output with an explicit cotangent seed.
5321 ///
5322 /// Like [`backward`](Self::backward), this rejects absent cotangents for
5323 /// reachable tracked leaves before accumulation and never returns an empty
5324 /// successful gradient bundle. [`EagerRuntime::vjp_optional`] exposes
5325 /// absent flow as `None` instead.
5326 ///
5327 /// This is the stateful eager VJP sugar: it returns the cotangent map and
5328 /// accumulates reachable tracked leaves into their `grad()` slots. Use
5329 /// [`EagerRuntime::vjp`] when the VJP result should be returned as a
5330 /// composable eager tensor without touching grad slots.
5331 ///
5332 /// # Examples
5333 ///
5334 /// ```
5335 /// use tenferro_ad::{EagerRuntime, EagerTensor, Tensor};
5336 ///
5337 /// let ctx = EagerRuntime::new()?;
5338 /// let x = EagerTensor::requires_grad_in(
5339 /// Tensor::from_vec_col_major(vec![2], vec![2.0_f64, 3.0]).unwrap(),
5340 /// ctx.clone(),
5341 /// )?;
5342 /// let seed = EagerTensor::from_tensor_in(
5343 /// Tensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0]).unwrap(),
5344 /// ctx,
5345 /// )?;
5346 /// let y = x.runtime().with_eager_session(&mut tenferro_cpu::CpuEngine::serial(), |s| s.mul(&x, &x))?;
5347 /// y.backward_with(&mut tenferro_cpu::CpuEngine::serial(), &seed)?;
5348 /// assert_eq!(x.grad()?.unwrap().as_slice::<f64>().unwrap(), &[4.0, 12.0]);
5349 /// # Ok::<(), tenferro_ad::Error>(())
5350 /// ```
5351 ///
5352 /// # Errors
5353 ///
5354 /// Returns [`Error::ContextMismatch`] when `cotangent` belongs to another
5355 /// eager runtime, [`Error::Validation`] when its shape or dtype is not a
5356 /// valid seed or a reachable tracked leaf has no cotangent,
5357 /// [`Error::RuntimeState`] when no tracked leaf produces a gradient,
5358 /// [`Error::UnsupportedAdRule`] for an unavailable reverse rule, or a typed
5359 /// backend/runtime-state error during execution.
5360 pub fn backward_with(
5361 &self,
5362 engine: &mut dyn ExecutionEngine,
5363 cotangent: &EagerTensor,
5364 ) -> Result<Gradients> {
5365 if !self.same_context(cotangent) {
5366 return Err(Error::ContextMismatch {
5367 lhs: self.ctx_id(),
5368 rhs: cotangent.ctx_id(),
5369 });
5370 }
5371 validate_seed_tensor("backward", self, cotangent)?;
5372 let seed = cotangent.to_tensor(&mut *engine)?;
5373 self.backward_from_seed(engine, seed)
5374 }
5375
5376 fn backward_from_seed(
5377 &self,
5378 engine: &mut dyn ExecutionEngine,
5379 seed: Tensor,
5380 ) -> Result<Gradients> {
5381 let cotangent =
5382 EagerTensor::new_result(Arc::clone(&self.ctx), eager_val_key(), seed, false, None)?;
5383 let candidate_keys = {
5384 let mut slots = self.ctx.lock_grad_slots()?;
5385 let mut keys = Vec::new();
5386 slots.retain(|key, slot| {
5387 if slot.upgrade().is_some() {
5388 keys.push(key.clone());
5389 true
5390 } else {
5391 false
5392 }
5393 });
5394 keys
5395 };
5396
5397 let mut targets = Vec::new();
5398 for key in candidate_keys {
5399 let Some(record) = self.ctx.value_record(&key)? else {
5400 continue;
5401 };
5402 if !record.requires_grad {
5403 continue;
5404 }
5405 // Intermediate and unrelated values also have slots; only request reachable leaves.
5406 let Some(input_key) = record
5407 .semantic_trace
5408 .as_ref()
5409 .and_then(TracedTensor::input_key)
5410 else {
5411 continue;
5412 };
5413 if !self
5414 .semantic_trace
5415 .as_ref()
5416 .is_some_and(|output| output.has_attached_input_key(&input_key))
5417 {
5418 continue;
5419 }
5420 targets.push((key, EagerTensor::from_record(record)));
5421 }
5422 let wrts = targets.iter().map(|(_, wrt)| wrt).collect::<Vec<_>>();
5423 let gradients = semantic_eager_vjp_many(&self.ctx, engine, self, &wrts, &cotangent)?;
5424 let cotangents = (|| {
5425 let session: &mut dyn Engine = &mut *engine;
5426 let mut cotangents = HashMap::new();
5427 for ((key, _), grad) in targets.into_iter().zip(gradients) {
5428 let grad = grad.ok_or_else(|| {
5429 Error::invalid_argument(
5430 "backward",
5431 ErrorPhase::GraphBuild,
5432 "wrt",
5433 format!(
5434 "no cotangent flow for reachable tracked leaf {key:?}; use EagerRuntime::vjp_optional to observe absent differentiation"
5435 ),
5436 )
5437 })?;
5438 let tensor = match grad.into_value() {
5439 Ok(tensor) => tensor,
5440 Err(IntoValueError::NotUnique(handle)) => {
5441 handle.duplicate_value_in_session(session)?
5442 }
5443 // A gradient whose retained layout is a view (for example a
5444 // transpose) cannot be extracted as an owned tensor; copy
5445 // it to a compact tensor instead.
5446 Err(IntoValueError::Extract { value, .. }) => {
5447 value.duplicate_value_in_session(session)?
5448 }
5449 };
5450 cotangents.insert(key, tensor);
5451 }
5452 self.ctx.store_grads(&cotangents, session)?;
5453 Ok::<_, Error>(cotangents)
5454 })()?;
5455 Gradients::from_tensors(cotangents)
5456 }
5457}
5458
5459/// Insert a weak registry entry, first dropping dead entries when the insert
5460/// would otherwise grow the table.
5461///
5462/// Dead entries are otherwise removed only when their own key is looked up, so
5463/// a long-running runtime would keep one entry per value it ever created.
5464/// Pruning at the growth point keeps the cost amortized O(1) per insert: a
5465/// sweep runs at most once per capacity doubling.
5466fn insert_pruning_dead<K: std::hash::Hash + Eq, V>(
5467 map: &mut HashMap<K, Weak<V>>,
5468 key: K,
5469 value: Weak<V>,
5470) {
5471 if map.len() == map.capacity() {
5472 map.retain(|_, entry| entry.strong_count() > 0);
5473 }
5474 map.insert(key, value);
5475}
5476
5477pub(crate) fn eager_val_key() -> ValueKey<StdTensorOp> {
5478 ValueKey::Input(next_input_key())
5479}
5480
5481pub(crate) struct RecordedEagerTrace {
5482 pub(crate) key: ValueKey<StdTensorOp>,
5483 pub(crate) trace: Option<EagerTrace>,
5484 pub(crate) requires_grad: bool,
5485}
5486
5487pub(crate) struct RecordedEagerOutputs {
5488 pub(crate) traces: Vec<RecordedEagerTrace>,
5489 pub(crate) semantic_traces: Vec<Option<TracedTensor>>,
5490}
5491
5492pub(crate) fn record_eager_outputs_in_session(
5493 op: &StdTensorOp,
5494 outputs: &[&Tensor],
5495 inputs: &[&EagerTensor],
5496 session: &mut dyn Engine,
5497) -> Result<RecordedEagerOutputs> {
5498 let metadata = outputs
5499 .iter()
5500 .map(|output| tensor_meta_from_tensor(output))
5501 .collect();
5502 record_eager_outputs_inner(op, metadata, inputs, Some(session))
5503}
5504
5505pub(crate) fn record_eager_value_outputs_in_session(
5506 op: &StdTensorOp,
5507 outputs: &[&TensorValue],
5508 inputs: &[&EagerTensor],
5509 session: &mut dyn Engine,
5510) -> Result<RecordedEagerOutputs> {
5511 let metadata = outputs
5512 .iter()
5513 .map(|output| tensor_meta_from_value(output))
5514 .collect();
5515 record_eager_outputs_inner(op, metadata, inputs, Some(session))
5516}
5517
5518fn record_eager_outputs_inner(
5519 op: &StdTensorOp,
5520 output_metadata: Vec<TensorMeta>,
5521 inputs: &[&EagerTensor],
5522 session: Option<&mut dyn Engine>,
5523) -> Result<RecordedEagerOutputs> {
5524 let semantic_traces = record_semantic_eager_outputs(op, &output_metadata, inputs, session)?;
5525 record_eager_outputs_from_metadata(output_metadata, semantic_traces, inputs)
5526}
5527
5528fn record_semantic_eager_outputs(
5529 op: &StdTensorOp,
5530 output_metadata: &[TensorMeta],
5531 inputs: &[&EagerTensor],
5532 mut session: Option<&mut dyn Engine>,
5533) -> Result<Vec<Option<TracedTensor>>> {
5534 // Materialize a constant semantic leaf for any untracked input that lost
5535 // its implicit semantic trace on the active-edge fast path. This keeps
5536 // "untracked constant feeds tracked AD" working (PyTorch-style: untracked
5537 // = constant leaf, no gradient flows to it) without re-recording every
5538 // untracked op at creation time.
5539 let mut owned_constants = Vec::<TracedTensor>::new();
5540 for input in inputs {
5541 if input.semantic_trace.is_none() {
5542 let value = match session.as_deref_mut() {
5543 Some(session) => input.duplicate_value_in_session(session)?,
5544 None => input.duplicate_host_value().ok_or_else(|| {
5545 Error::unsupported(
5546 "record_semantic_eager_outputs",
5547 ErrorPhase::Execution,
5548 "a non-host untracked input needs an engine to become a constant leaf",
5549 )
5550 })?,
5551 };
5552 owned_constants.push(TracedTensor::from_tensor_symbolic_shape(value)?);
5553 }
5554 }
5555 let mut constants = owned_constants.iter();
5556 let mut semantic_inputs: Vec<&TracedTensor> = inputs
5557 .iter()
5558 .map(|input| {
5559 input
5560 .semantic_trace
5561 .as_ref()
5562 .unwrap_or_else(|| constants.next().expect("materialized constant"))
5563 })
5564 .collect();
5565 let promotion_plan =
5566 eager_input_promotion_plan(op, inputs.len(), |index| inputs[index].dtype());
5567 // Mirror eager execution in the deferred carrier only. The concrete
5568 // tensors have already been promoted at the execution boundary, so these
5569 // casts add semantic graph nodes without an eager copy or backend kernel.
5570 let promoted_semantic_inputs = if semantic_inputs.iter().enumerate().any(|(index, semantic)| {
5571 semantic.dtype != promotion_plan.target_dtype(index, semantic.dtype)
5572 }) {
5573 Some(
5574 semantic_inputs
5575 .iter()
5576 .enumerate()
5577 .map(|(index, &semantic)| {
5578 let target = promotion_plan.target_dtype(index, semantic.dtype);
5579 if semantic.dtype == target {
5580 Ok(Cow::Borrowed(semantic))
5581 } else {
5582 semantic.cast(target).map(Cow::Owned)
5583 }
5584 })
5585 .collect::<Result<Vec<Cow<'_, TracedTensor>>>>()?,
5586 )
5587 } else {
5588 None
5589 };
5590 if let Some(promoted_semantic_inputs) = &promoted_semantic_inputs {
5591 semantic_inputs = promoted_semantic_inputs.iter().map(Cow::as_ref).collect();
5592 }
5593 let exact_semantic_inputs = if matches!(op, StdTensorOp::Concatenate { .. }) {
5594 Some(
5595 semantic_inputs
5596 .iter()
5597 .zip(inputs)
5598 .map(|(&semantic, input)| {
5599 if semantic.is_concrete_shape() {
5600 Ok(semantic.clone())
5601 } else {
5602 semantic.reshape(input.shape())
5603 }
5604 })
5605 .collect::<Result<Vec<_>>>()?,
5606 )
5607 } else {
5608 None
5609 };
5610 if let Some(exact_semantic_inputs) = &exact_semantic_inputs {
5611 semantic_inputs = exact_semantic_inputs.iter().collect();
5612 }
5613 // Deferred materialization (issue #1665 steps 6-7): append only a raw
5614 // carrier. The runtime helper retains metadata scopes introduced by the
5615 // promotion/exactification helpers without analyzing this operation.
5616 let outputs = tenferro_runtime::extension::append_raw_eager_outputs(
5617 op.clone(),
5618 &semantic_inputs,
5619 output_metadata,
5620 )?;
5621 Ok(outputs.into_iter().map(Some).collect())
5622}
5623
5624#[cfg(test)]
5625fn record_eager_graph_outputs(
5626 graph: &Graph<StdTensorOp>,
5627 graph_input_keys: &[TensorInputKey],
5628 outputs: &[Tensor],
5629 inputs: &[&EagerTensor],
5630) -> Result<RecordedEagerOutputs> {
5631 let semantic_traces = record_semantic_eager_graph_outputs(graph, graph_input_keys, inputs)?;
5632 let output_metadata = outputs.iter().map(tensor_meta_from_tensor);
5633 record_eager_outputs_from_metadata(output_metadata, semantic_traces, inputs)
5634}
5635
5636#[cfg(test)]
5637fn record_semantic_eager_graph_outputs(
5638 graph: &Graph<StdTensorOp>,
5639 graph_input_keys: &[TensorInputKey],
5640 inputs: &[&EagerTensor],
5641) -> Result<Vec<Option<TracedTensor>>> {
5642 let Some(semantic_inputs) = inputs
5643 .iter()
5644 .map(|input| input.semantic_trace.as_ref())
5645 .collect::<Option<Vec<_>>>()
5646 else {
5647 return Ok(vec![None; graph.outputs().len()]);
5648 };
5649 if graph_input_keys.len() != semantic_inputs.len() {
5650 return Err(Error::Internal(format!(
5651 "semantic graph recording expected {} input keys, got {}",
5652 semantic_inputs.len(),
5653 graph_input_keys.len()
5654 )));
5655 }
5656
5657 let mut values = HashMap::new();
5658 for (key, tensor) in graph_input_keys.iter().zip(semantic_inputs) {
5659 values.insert(ValueKey::Input(key.clone()), tensor.clone());
5660 }
5661
5662 for op_node in graph.operations() {
5663 let input_values = op_node
5664 .inputs
5665 .iter()
5666 .map(|input| {
5667 let key = match input {
5668 ValueRef::Local(local_id) => &graph.values()[*local_id].key,
5669 ValueRef::External(key) => key,
5670 };
5671 values.get(key).cloned().ok_or_else(|| {
5672 Error::Internal(format!(
5673 "semantic graph recording missing value for {key:?}"
5674 ))
5675 })
5676 })
5677 .collect::<Result<Vec<_>>>()?;
5678 let input_refs = input_values.iter().collect::<Vec<_>>();
5679 let semantic_outputs = match &op_node.operation {
5680 StdTensorOp::Extension(ext) => {
5681 tenferro_runtime::extension::apply(Arc::clone(ext), &input_refs)?
5682 }
5683 op => tenferro_runtime::extension::apply_standard_op(op.clone(), &input_refs)?,
5684 };
5685 if semantic_outputs.len() != op_node.outputs.len() {
5686 return Err(Error::Internal(format!(
5687 "semantic graph recording expected {} outputs for {:?}, got {}",
5688 op_node.outputs.len(),
5689 op_node.operation,
5690 semantic_outputs.len()
5691 )));
5692 }
5693 for (output_id, output) in op_node.outputs.iter().copied().zip(semantic_outputs) {
5694 values.insert(graph.values()[output_id].key.clone(), output);
5695 }
5696 }
5697
5698 graph
5699 .outputs()
5700 .iter()
5701 .map(|&output_id| {
5702 let key = &graph.values()[output_id].key;
5703 values.get(key).cloned().map(Some).ok_or_else(|| {
5704 Error::Internal(format!(
5705 "semantic graph recording missing output for {key:?}"
5706 ))
5707 })
5708 })
5709 .collect()
5710}
5711
5712fn record_eager_outputs_from_metadata(
5713 output_metadata: impl IntoIterator<Item = TensorMeta>,
5714 semantic_traces: Vec<Option<TracedTensor>>,
5715 inputs: &[&EagerTensor],
5716) -> Result<RecordedEagerOutputs> {
5717 let output_metadata = output_metadata.into_iter().collect::<Vec<_>>();
5718 if semantic_traces.len() != output_metadata.len() {
5719 return Err(Error::Internal(format!(
5720 "eager recording expected {} semantic traces, got {}",
5721 output_metadata.len(),
5722 semantic_traces.len()
5723 )));
5724 }
5725 let requires_grad =
5726 eager_grad_recording_enabled() && inputs.iter().any(|input| input.requires_grad);
5727 let trace_count = output_metadata.len();
5728 let residual_trace = EagerTrace::new(inputs);
5729 let traces = (0..trace_count)
5730 .map(|_| RecordedEagerTrace {
5731 key: eager_val_key(),
5732 trace: Some(residual_trace.clone()),
5733 requires_grad,
5734 })
5735 .collect();
5736
5737 Ok(RecordedEagerOutputs {
5738 traces,
5739 semantic_traces,
5740 })
5741}
5742
5743fn tensor_meta_from_value(value: &TensorValue) -> TensorMeta {
5744 TensorMeta::exact(
5745 value.dtype(),
5746 value.shape().iter().copied().map(SymDim::from).collect(),
5747 )
5748}
5749
5750#[cfg(test)]
5751pub(crate) fn zero_like_tensor(input: &Tensor, backend: &mut dyn Engine) -> Result<Tensor> {
5752 let host = match input.dtype() {
5753 // A caller-owned payload has no zero-like runtime tensor.
5754 DType::External(type_id) => {
5755 return Err(Error::unsupported(
5756 "zero_like_tensor",
5757 ErrorPhase::GraphBuild,
5758 format!(
5759 "an externally defined payload ({:?}) has no zero-like runtime tensor",
5760 DType::External(type_id)
5761 ),
5762 ));
5763 }
5764 DType::F32 => Tensor::from_typed::<f32>(TypedTensor::zeros(input.shape().to_vec())?),
5765 DType::F64 => Tensor::from_typed::<f64>(TypedTensor::zeros(input.shape().to_vec())?),
5766 DType::I32 => Tensor::from_typed::<i32>(TypedTensor::zeros(input.shape().to_vec())?),
5767 DType::I64 => Tensor::from_typed::<i64>(TypedTensor::zeros(input.shape().to_vec())?),
5768 DType::Bool => Tensor::from_typed::<bool>(TypedTensor::from_vec_col_major(
5769 input.shape().to_vec(),
5770 vec![false; input.shape().iter().product()],
5771 )?),
5772 DType::C32 => Tensor::from_typed::<tenferro_tensor::Complex32>(TypedTensor::zeros(
5773 input.shape().to_vec(),
5774 )?),
5775 DType::C64 => Tensor::from_typed::<tenferro_tensor::Complex64>(TypedTensor::zeros(
5776 input.shape().to_vec(),
5777 )?),
5778 };
5779 backend
5780 .upload_host_tensor(TensorRead::from_tensor(&host))
5781 .map_err(Error::from)
5782}
5783
5784pub(crate) fn one_like_tensor(input: &Tensor, session: &mut dyn Engine) -> Result<Tensor> {
5785 let host = ones_tensor(input.dtype(), input.shape().to_vec())?;
5786 session
5787 .upload_host_tensor(TensorRead::from_tensor(&host))
5788 .map_err(Error::from)
5789}
5790
5791#[cfg(test)]
5792mod tests;