tenferro_ad/context.rs
1//! Explicit ownership for automatic-differentiation rule sets.
2
3use std::sync::Arc;
4
5use tenferro_runtime::program::FrozenProgram;
6use tenferro_runtime::{CacheStats, Result, TracedTensor};
7
8// SemanticCompatDispatcher removed in Unification 7.
9// Extension AD is handled exclusively by SemanticExtensionRuleSet.
10use crate::semantic_extension::{SemanticExtensionRegistryError, SemanticExtensionRuleSet};
11use crate::semantic_transform::{
12 semantic_jvp, semantic_vjp, SemanticAdProgram, SemanticAdTransformError,
13};
14use crate::transform_cache::{
15 AdTransformCache, AdTransformCacheLimits, SemanticAdTransformCacheKey,
16};
17
18/// Stats for caches owned by an [`AdContext`].
19///
20/// `retained_bytes` fields are logical payload estimates, not process RSS.
21///
22/// # Examples
23///
24/// ```rust
25/// use tenferro_ad::AdContext;
26///
27/// let ad = AdContext::builder().build().unwrap();
28/// assert_eq!(ad.cache_stats().unwrap().ad_transforms.entries, 0);
29/// ```
30#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
31pub struct AdContextCacheStats {
32 /// AD transform graph memoization cache.
33 pub ad_transforms: CacheStats,
34}
35
36/// Explicit automatic-differentiation context.
37///
38/// `AdContext` owns the extension AD rules used by traced AD transforms.
39/// It also owns the AD transform cache shared by context-driven traced AD and
40/// eager runtimes created from this context.
41///
42/// # Examples
43///
44/// ```rust
45/// use tenferro_ad::AdContext;
46///
47/// let ad = AdContext::builder().build().unwrap();
48/// assert!(ad
49/// .semantic_extension_rules()
50/// .lookup_linearize("example.missing.v1")
51/// .is_none());
52/// ```
53#[derive(Clone, Debug)]
54pub struct AdContext {
55 semantic_extension_rules: SemanticExtensionRuleSet,
56 ad_transform_cache: Arc<AdTransformCache>,
57}
58
59impl AdContext {
60 /// Start building an explicit AD context.
61 ///
62 /// # Examples
63 ///
64 /// ```rust
65 /// use tenferro_ad::AdContext;
66 ///
67 /// let _builder = AdContext::builder();
68 /// ```
69 pub fn builder() -> AdContextBuilder {
70 AdContextBuilder::default()
71 }
72
73 pub(crate) fn with_rules_and_transform_cache(
74 semantic_extension_rules: SemanticExtensionRuleSet,
75 ad_transform_cache: Arc<AdTransformCache>,
76 ) -> Self {
77 Self {
78 semantic_extension_rules,
79 ad_transform_cache,
80 }
81 }
82
83 /// Return semantic-program extension AD rules owned by this context.
84 ///
85 /// # Examples
86 ///
87 /// ```rust
88 /// use tenferro_ad::AdContext;
89 ///
90 /// let ad = AdContext::builder().build().unwrap();
91 /// assert!(ad
92 /// .semantic_extension_rules()
93 /// .lookup_linearize("example.missing.v1")
94 /// .is_none());
95 /// ```
96 pub fn semantic_extension_rules(&self) -> &SemanticExtensionRuleSet {
97 &self.semantic_extension_rules
98 }
99
100 /// Transform a frozen semantic program into its forward-mode derivative.
101 ///
102 /// `active_inputs` follows source-program input order. Active tangent
103 /// seeds are appended after all primal inputs.
104 ///
105 /// # Errors
106 ///
107 /// Returns [`SemanticAdTransformError::ActivityArity`] when
108 /// `active_inputs` has the wrong length,
109 /// [`SemanticAdTransformError::Extension`] when an extension rule rejects
110 /// the transform, or the corresponding `Query`, `Build`, `Finish`, or
111 /// `Cache` variant when program import, construction, finalization, or
112 /// cache access fails.
113 pub fn jvp_program(
114 &self,
115 input: &FrozenProgram,
116 active_inputs: &[bool],
117 ) -> std::result::Result<SemanticAdProgram, SemanticAdTransformError> {
118 let key = SemanticAdTransformCacheKey::jvp(input, active_inputs);
119 if let Some(cached) = self
120 .ad_transform_cache
121 .get_semantic(&key, input)
122 .map_err(SemanticAdTransformError::Cache)?
123 {
124 return cached
125 .as_ref()
126 .with_input_prefix_bindings_from(input)
127 .map_err(SemanticAdTransformError::from);
128 }
129 let transformed = semantic_jvp(input, active_inputs, &self.semantic_extension_rules)?;
130 self.ad_transform_cache
131 .put_semantic(key, input, Arc::new(transformed.clone()))
132 .map_err(SemanticAdTransformError::Cache)?;
133 Ok(transformed)
134 }
135
136 /// Transform a frozen semantic program into its reverse-mode derivative.
137 ///
138 /// `active_inputs` selects requested primal-input cotangents and
139 /// `active_outputs` selects primal outputs that receive appended seeds.
140 ///
141 /// # Errors
142 ///
143 /// Returns [`SemanticAdTransformError::ActivityArity`] when either activity
144 /// mask has the wrong length,
145 /// [`SemanticAdTransformError::Extension`] when an extension rule rejects
146 /// the transform, or the corresponding `Query`, `Build`, `Finish`, or
147 /// `Cache` variant when program import, construction, finalization, or
148 /// cache access fails.
149 pub fn vjp_program(
150 &self,
151 input: &FrozenProgram,
152 active_inputs: &[bool],
153 active_outputs: &[bool],
154 ) -> std::result::Result<SemanticAdProgram, SemanticAdTransformError> {
155 let key = SemanticAdTransformCacheKey::vjp(input, active_inputs, active_outputs);
156 if let Some(cached) = self
157 .ad_transform_cache
158 .get_semantic(&key, input)
159 .map_err(SemanticAdTransformError::Cache)?
160 {
161 return cached
162 .as_ref()
163 .with_input_prefix_bindings_from(input)
164 .map_err(SemanticAdTransformError::from);
165 }
166 let transformed = semantic_vjp(
167 input,
168 active_inputs,
169 active_outputs,
170 &self.semantic_extension_rules,
171 )?;
172 self.ad_transform_cache
173 .put_semantic(key, input, Arc::new(transformed.clone()))
174 .map_err(SemanticAdTransformError::Cache)?;
175 Ok(transformed)
176 }
177
178 pub(crate) fn ad_transform_cache(&self) -> Arc<AdTransformCache> {
179 Arc::clone(&self.ad_transform_cache)
180 }
181
182 /// Return AD transform cache retention limits.
183 ///
184 /// # Examples
185 ///
186 /// ```rust
187 /// use tenferro_ad::AdContext;
188 ///
189 /// let ad = AdContext::builder().build().unwrap();
190 /// assert!(ad.ad_transform_cache_limits().unwrap().max_entries().get() > 0);
191 /// ```
192 ///
193 /// # Errors
194 ///
195 /// Returns [`tenferro_runtime::Error::RuntimeState`] if the cache lock is
196 /// poisoned or its state cannot be inspected.
197 pub fn ad_transform_cache_limits(&self) -> Result<AdTransformCacheLimits> {
198 self.ad_transform_cache.limits()
199 }
200
201 /// Replace AD transform cache retention limits.
202 ///
203 /// # Examples
204 ///
205 /// ```rust
206 /// use std::num::NonZeroUsize;
207 /// use tenferro_ad::{AdContext, AdTransformCacheLimits};
208 ///
209 /// let ad = AdContext::builder().build().unwrap();
210 /// let limits = AdTransformCacheLimits::new(NonZeroUsize::new(1).unwrap());
211 /// ad.set_ad_transform_cache_limits(limits).unwrap();
212 /// assert_eq!(ad.ad_transform_cache_limits().unwrap(), limits);
213 /// ```
214 ///
215 /// # Errors
216 ///
217 /// Returns [`tenferro_runtime::Error::RuntimeState`] if the cache lock is
218 /// poisoned while updating the limits.
219 pub fn set_ad_transform_cache_limits(&self, limits: AdTransformCacheLimits) -> Result<()> {
220 self.ad_transform_cache.set_limits(limits)
221 }
222
223 /// Clear AD transform cache entries owned by this context.
224 ///
225 /// # Examples
226 ///
227 /// ```rust
228 /// use tenferro_ad::AdContext;
229 ///
230 /// let ad = AdContext::builder().build().unwrap();
231 /// ad.clear_ad_transform_caches().unwrap();
232 /// assert_eq!(ad.ad_transform_cache_stats().unwrap().entries, 0);
233 /// ```
234 ///
235 /// # Errors
236 ///
237 /// Returns [`tenferro_runtime::Error::RuntimeState`] if the cache lock is
238 /// poisoned while clearing entries.
239 pub fn clear_ad_transform_caches(&self) -> Result<()> {
240 self.ad_transform_cache.clear()
241 }
242
243 /// Return AD transform cache-entry and retained-byte stats.
244 ///
245 /// # Examples
246 ///
247 /// ```rust
248 /// use tenferro_ad::AdContext;
249 ///
250 /// let ad = AdContext::builder().build().unwrap();
251 /// assert_eq!(ad.ad_transform_cache_stats().unwrap().entries, 0);
252 /// ```
253 ///
254 /// # Errors
255 ///
256 /// Returns [`tenferro_runtime::Error::RuntimeState`] if the cache lock is
257 /// poisoned while collecting statistics.
258 pub fn ad_transform_cache_stats(&self) -> Result<CacheStats> {
259 self.ad_transform_cache.stats()
260 }
261
262 /// Clear every cache owned by this AD context.
263 ///
264 /// # Examples
265 ///
266 /// ```rust
267 /// use tenferro_ad::AdContext;
268 ///
269 /// let ad = AdContext::builder().build().unwrap();
270 /// ad.clear_caches().unwrap();
271 /// assert_eq!(ad.cache_stats().unwrap().ad_transforms.entries, 0);
272 /// ```
273 ///
274 /// # Errors
275 ///
276 /// Returns [`tenferro_runtime::Error::RuntimeState`] if either owned cache
277 /// cannot be locked because its state is poisoned.
278 pub fn clear_caches(&self) -> Result<()> {
279 self.clear_ad_transform_caches()
280 }
281
282 /// Return aggregate cache-entry and retained-byte stats for this AD context.
283 ///
284 /// # Examples
285 ///
286 /// ```rust
287 /// use tenferro_ad::AdContext;
288 ///
289 /// let ad = AdContext::builder().build().unwrap();
290 /// assert_eq!(ad.cache_stats().unwrap().ad_transforms.retained_bytes, 0);
291 /// ```
292 ///
293 /// # Errors
294 ///
295 /// Returns [`tenferro_runtime::Error::RuntimeState`] if an owned cache lock
296 /// is poisoned while collecting statistics.
297 pub fn cache_stats(&self) -> Result<AdContextCacheStats> {
298 Ok(AdContextCacheStats {
299 ad_transforms: self.ad_transform_cache_stats()?,
300 })
301 }
302
303 /// Gradient of a scalar traced output with respect to a traced input.
304 ///
305 /// For complex scalar outputs, tenferro returns the Hermitian-adjoint
306 /// cotangent. To compare seed-`1` scalar gradients with JAX's public
307 /// `grad` values, use the complex conjugate of this result. See
308 /// <https://tensor4all.org/tenferro-rs/guides/complex-ad.html>.
309 ///
310 /// # Examples
311 ///
312 /// ```rust
313 /// use tenferro_ad::AdContext;
314 /// use tenferro_runtime::TracedTensor;
315 ///
316 /// let ad = AdContext::builder().build().unwrap();
317 /// let x = TracedTensor::from_vec_col_major(vec![], vec![3.0_f64]).unwrap();
318 /// let loss = (&x * &x).unwrap();
319 /// let grad = ad.grad(&loss, &x).unwrap();
320 /// assert_eq!(grad.rank, 0);
321 /// ```
322 ///
323 /// # Errors
324 ///
325 /// Returns [`tenferro_runtime::Error::NonScalarGrad`] when `output` is not
326 /// scalar, [`tenferro_runtime::Error::UnsupportedAdRule`] when a graph op
327 /// lacks a registered rule, or a typed [`tenferro_runtime::Error::Validation`]
328 /// / backend error when graph metadata or execution is invalid.
329 pub fn grad(&self, output: &TracedTensor, wrt: &TracedTensor) -> Result<TracedTensor> {
330 crate::traced::grad_with_rules_and_cache(
331 output,
332 wrt,
333 &self.semantic_extension_rules,
334 Some(self.ad_transform_cache.as_ref()),
335 )
336 }
337
338 /// Gradient that returns `None` when `wrt` is inactive.
339 ///
340 /// # Examples
341 ///
342 /// ```rust
343 /// use tenferro_ad::AdContext;
344 /// use tenferro_runtime::TracedTensor;
345 ///
346 /// let ad = AdContext::builder().build().unwrap();
347 /// let x = TracedTensor::from_vec_col_major(vec![], vec![3.0_f64]).unwrap();
348 /// let loss = (&x * &x).unwrap();
349 /// assert!(ad.grad_optional(&loss, &x).unwrap().is_some());
350 /// ```
351 ///
352 /// # Errors
353 ///
354 /// Returns [`tenferro_runtime::Error::NonScalarGrad`] for a non-scalar
355 /// output, [`tenferro_runtime::Error::UnsupportedAdRule`] for an
356 /// unregistered AD rule, or a typed [`tenferro_runtime::Error::Validation`]
357 /// / backend error from graph construction and
358 /// execution.
359 pub fn grad_optional(
360 &self,
361 output: &TracedTensor,
362 wrt: &TracedTensor,
363 ) -> Result<Option<TracedTensor>> {
364 crate::traced::grad_optional_with_rules_and_cache(
365 output,
366 wrt,
367 &self.semantic_extension_rules,
368 Some(self.ad_transform_cache.as_ref()),
369 )
370 }
371
372 /// Forward-mode Jacobian-vector product.
373 ///
374 /// # Examples
375 ///
376 /// ```rust
377 /// use tenferro_ad::AdContext;
378 /// use tenferro_runtime::TracedTensor;
379 ///
380 /// let ad = AdContext::builder().build().unwrap();
381 /// let x = TracedTensor::from_vec_col_major(vec![], vec![3.0_f64]).unwrap();
382 /// let dx = TracedTensor::from_vec_col_major(vec![], vec![1.0_f64]).unwrap();
383 /// let y = (&x * &x).unwrap();
384 /// let dy = ad.jvp(&y, &x, &dx).unwrap();
385 /// assert_eq!(dy.rank, 0);
386 /// ```
387 ///
388 /// # Errors
389 ///
390 /// Returns [`tenferro_runtime::Error::UnsupportedAdRule`] when the graph
391 /// has no JVP rule, [`tenferro_runtime::Error::Validation`] for
392 /// inconsistent tangent metadata, or a typed backend/runtime-state error
393 /// during evaluation.
394 pub fn jvp(
395 &self,
396 output: &TracedTensor,
397 wrt: &TracedTensor,
398 tangent: &TracedTensor,
399 ) -> Result<TracedTensor> {
400 crate::traced::jvp_with_rules_and_cache(
401 output,
402 wrt,
403 tangent,
404 &self.semantic_extension_rules,
405 Some(self.ad_transform_cache.as_ref()),
406 )
407 }
408
409 /// Forward-mode Jacobian-vector product that returns `None` for inactive output.
410 ///
411 /// # Examples
412 ///
413 /// ```rust
414 /// use tenferro_ad::AdContext;
415 /// use tenferro_runtime::TracedTensor;
416 ///
417 /// let ad = AdContext::builder().build().unwrap();
418 /// let x = TracedTensor::from_vec_col_major(vec![], vec![3.0_f64]).unwrap();
419 /// let dx = TracedTensor::from_vec_col_major(vec![], vec![1.0_f64]).unwrap();
420 /// let y = (&x * &x).unwrap();
421 /// assert!(ad.jvp_optional(&y, &x, &dx).unwrap().is_some());
422 /// ```
423 ///
424 /// # Errors
425 ///
426 /// Returns [`tenferro_runtime::Error::UnsupportedAdRule`] when the graph
427 /// has no JVP rule, [`tenferro_runtime::Error::Validation`] for
428 /// inconsistent tangent metadata, or a typed backend/runtime-state error
429 /// during evaluation.
430 pub fn jvp_optional(
431 &self,
432 output: &TracedTensor,
433 wrt: &TracedTensor,
434 tangent: &TracedTensor,
435 ) -> Result<Option<TracedTensor>> {
436 crate::traced::jvp_optional_with_rules_and_cache(
437 output,
438 wrt,
439 tangent,
440 &self.semantic_extension_rules,
441 Some(self.ad_transform_cache.as_ref()),
442 )
443 }
444
445 /// Reverse-mode vector-Jacobian product.
446 ///
447 /// Complex cotangents use tenferro's Hermitian real-inner-product
448 /// convention. Non-real complex cotangent seeds therefore need an explicit
449 /// seed-convention comparison when matching JAX. See
450 /// <https://tensor4all.org/tenferro-rs/guides/complex-ad.html>.
451 ///
452 /// # Examples
453 ///
454 /// ```rust
455 /// use tenferro_ad::AdContext;
456 /// use tenferro_runtime::TracedTensor;
457 ///
458 /// let ad = AdContext::builder().build().unwrap();
459 /// let x = TracedTensor::from_vec_col_major(vec![], vec![3.0_f64]).unwrap();
460 /// let dy = TracedTensor::from_vec_col_major(vec![], vec![1.0_f64]).unwrap();
461 /// let y = (&x * &x).unwrap();
462 /// let dx = ad.vjp(&y, &x, &dy).unwrap();
463 /// assert_eq!(dx.rank, 0);
464 /// ```
465 ///
466 /// # Errors
467 ///
468 /// Returns [`tenferro_runtime::Error::Validation`] when the cotangent
469 /// metadata is incompatible, [`tenferro_runtime::Error::UnsupportedAdRule`]
470 /// when a VJP rule is unavailable, or a typed backend/runtime-state error
471 /// during execution.
472 pub fn vjp(
473 &self,
474 output: &TracedTensor,
475 wrt: &TracedTensor,
476 cotangent: &TracedTensor,
477 ) -> Result<TracedTensor> {
478 crate::traced::vjp_with_rules_and_cache(
479 output,
480 wrt,
481 cotangent,
482 &self.semantic_extension_rules,
483 Some(self.ad_transform_cache.as_ref()),
484 )
485 }
486
487 /// Reverse-mode vector-Jacobian product that returns `None` for inactive input.
488 ///
489 /// # Examples
490 ///
491 /// ```rust
492 /// use tenferro_ad::AdContext;
493 /// use tenferro_runtime::TracedTensor;
494 ///
495 /// let ad = AdContext::builder().build().unwrap();
496 /// let x = TracedTensor::from_vec_col_major(vec![], vec![3.0_f64]).unwrap();
497 /// let dy = TracedTensor::from_vec_col_major(vec![], vec![1.0_f64]).unwrap();
498 /// let y = (&x * &x).unwrap();
499 /// assert!(ad.vjp_optional(&y, &x, &dy).unwrap().is_some());
500 /// ```
501 ///
502 /// # Errors
503 ///
504 /// Returns [`tenferro_runtime::Error::Validation`] when the cotangent
505 /// metadata is incompatible, [`tenferro_runtime::Error::UnsupportedAdRule`]
506 /// when a VJP rule is unavailable, or a typed backend/runtime-state error
507 /// during execution.
508 pub fn vjp_optional(
509 &self,
510 output: &TracedTensor,
511 wrt: &TracedTensor,
512 cotangent: &TracedTensor,
513 ) -> Result<Option<TracedTensor>> {
514 crate::traced::vjp_optional_with_rules_and_cache(
515 output,
516 wrt,
517 cotangent,
518 &self.semantic_extension_rules,
519 Some(self.ad_transform_cache.as_ref()),
520 )
521 }
522}
523
524/// Builder for [`AdContext`].
525///
526/// # Examples
527///
528/// ```rust
529/// use tenferro_ad::AdContextBuilder;
530///
531/// let ad = AdContextBuilder::new().build().unwrap();
532/// assert!(ad
533/// .semantic_extension_rules()
534/// .lookup_linearize("example.missing.v1")
535/// .is_none());
536/// ```
537#[derive(Clone, Debug, Default)]
538pub struct AdContextBuilder {
539 semantic_extension_rules: SemanticExtensionRuleSet,
540}
541
542impl AdContextBuilder {
543 /// Create an empty builder.
544 ///
545 /// # Examples
546 ///
547 /// ```rust
548 /// use tenferro_ad::AdContextBuilder;
549 ///
550 /// let _builder = AdContextBuilder::new();
551 /// ```
552 pub fn new() -> Self {
553 Self::default()
554 }
555
556 /// Include an owned semantic-program extension AD rule set.
557 ///
558 /// # Errors
559 ///
560 /// Returns [`SemanticExtensionRegistryError::MalformedFamilyId`] when a
561 /// family identifier is invalid, or
562 /// [`SemanticExtensionRegistryError::DuplicateRule`] when the same family
563 /// and role were already supplied.
564 pub fn with_semantic_extension_rules(
565 mut self,
566 rules: SemanticExtensionRuleSet,
567 ) -> std::result::Result<Self, SemanticExtensionRegistryError> {
568 self.semantic_extension_rules.merge(rules)?;
569 Ok(self)
570 }
571
572 /// Build the context.
573 ///
574 /// Semantic extension rules have already been validated and merged by
575 /// [`Self::with_semantic_extension_rules`].
576 ///
577 /// # Examples
578 ///
579 /// ```rust
580 /// use tenferro_ad::AdContext;
581 ///
582 /// let ad = AdContext::builder().build().unwrap();
583 /// assert!(ad
584 /// .semantic_extension_rules()
585 /// .lookup_linearize("example.missing.v1")
586 /// .is_none());
587 /// ```
588 ///
589 /// # Errors
590 ///
591 /// The error type is [`std::convert::Infallible`], so this finalization step
592 /// never returns `Err` after semantic rule registration. It retains a
593 /// `Result` so callers can compose it with the fallible registration step.
594 pub fn build(self) -> std::result::Result<AdContext, std::convert::Infallible> {
595 Ok(AdContext {
596 semantic_extension_rules: self.semantic_extension_rules,
597 ad_transform_cache: Arc::new(AdTransformCache::new()),
598 })
599 }
600}