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