tenferro_ad/context.rs
1//! Explicit ownership for automatic-differentiation rule sets.
2
3use std::sync::Arc;
4
5use tenferro_ops::{ExtensionRegistryError, ExtensionRuleSet};
6use tenferro_runtime::{CacheStats, Result, TracedTensor};
7
8use crate::transform_cache::{AdTransformCache, AdTransformCacheLimits};
9
10/// Stats for caches owned by an [`AdContext`].
11///
12/// `retained_bytes` fields are logical payload estimates, not process RSS.
13///
14/// # Examples
15///
16/// ```rust
17/// use tenferro_ad::AdContext;
18///
19/// let ad = AdContext::builder().build().unwrap();
20/// assert_eq!(ad.cache_stats().unwrap().ad_transforms.entries, 0);
21/// ```
22#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
23pub struct AdContextCacheStats {
24 /// AD transform graph memoization cache.
25 pub ad_transforms: CacheStats,
26}
27
28/// Explicit automatic-differentiation context.
29///
30/// `AdContext` owns the extension AD rules used by traced AD transforms.
31/// It also owns the AD transform cache shared by context-driven traced AD and
32/// eager runtimes created from this context.
33///
34/// # Examples
35///
36/// ```rust
37/// use tenferro_ad::AdContext;
38///
39/// let ad = AdContext::builder().build().unwrap();
40/// assert!(ad.extension_rules().lookup_linearize("example.missing.v1").is_none());
41/// ```
42#[derive(Clone, Debug)]
43pub struct AdContext {
44 extension_rules: ExtensionRuleSet,
45 ad_transform_cache: Arc<AdTransformCache>,
46}
47
48impl AdContext {
49 /// Start building an explicit AD context.
50 ///
51 /// # Examples
52 ///
53 /// ```rust
54 /// use tenferro_ad::AdContext;
55 ///
56 /// let _builder = AdContext::builder();
57 /// ```
58 pub fn builder() -> AdContextBuilder {
59 AdContextBuilder::default()
60 }
61
62 /// Return the extension rules owned by this context.
63 ///
64 /// # Examples
65 ///
66 /// ```rust
67 /// use tenferro_ad::AdContext;
68 ///
69 /// let ad = AdContext::builder().build().unwrap();
70 /// assert!(!ad.extension_rules().is_linearize_registered("example.missing.v1"));
71 /// ```
72 pub fn extension_rules(&self) -> &ExtensionRuleSet {
73 &self.extension_rules
74 }
75
76 pub(crate) fn extension_rule_set(&self) -> ExtensionRuleSet {
77 self.extension_rules.clone()
78 }
79
80 pub(crate) fn ad_transform_cache(&self) -> Arc<AdTransformCache> {
81 Arc::clone(&self.ad_transform_cache)
82 }
83
84 /// Return AD transform cache retention limits.
85 ///
86 /// # Examples
87 ///
88 /// ```rust
89 /// use tenferro_ad::AdContext;
90 ///
91 /// let ad = AdContext::builder().build().unwrap();
92 /// assert!(ad.ad_transform_cache_limits().unwrap().max_entries().get() > 0);
93 /// ```
94 pub fn ad_transform_cache_limits(&self) -> Result<AdTransformCacheLimits> {
95 self.ad_transform_cache.limits()
96 }
97
98 /// Replace AD transform cache retention limits.
99 ///
100 /// # Examples
101 ///
102 /// ```rust
103 /// use std::num::NonZeroUsize;
104 /// use tenferro_ad::{AdContext, AdTransformCacheLimits};
105 ///
106 /// let ad = AdContext::builder().build().unwrap();
107 /// let limits = AdTransformCacheLimits::new(NonZeroUsize::new(1).unwrap());
108 /// ad.set_ad_transform_cache_limits(limits).unwrap();
109 /// assert_eq!(ad.ad_transform_cache_limits().unwrap(), limits);
110 /// ```
111 pub fn set_ad_transform_cache_limits(&self, limits: AdTransformCacheLimits) -> Result<()> {
112 self.ad_transform_cache.set_limits(limits)
113 }
114
115 /// Clear AD transform cache entries owned by this context.
116 ///
117 /// # Examples
118 ///
119 /// ```rust
120 /// use tenferro_ad::AdContext;
121 ///
122 /// let ad = AdContext::builder().build().unwrap();
123 /// ad.clear_ad_transform_caches().unwrap();
124 /// assert_eq!(ad.ad_transform_cache_stats().unwrap().entries, 0);
125 /// ```
126 pub fn clear_ad_transform_caches(&self) -> Result<()> {
127 self.ad_transform_cache.clear()
128 }
129
130 /// Return AD transform cache-entry and retained-byte stats.
131 ///
132 /// # Examples
133 ///
134 /// ```rust
135 /// use tenferro_ad::AdContext;
136 ///
137 /// let ad = AdContext::builder().build().unwrap();
138 /// assert_eq!(ad.ad_transform_cache_stats().unwrap().entries, 0);
139 /// ```
140 pub fn ad_transform_cache_stats(&self) -> Result<CacheStats> {
141 self.ad_transform_cache.stats()
142 }
143
144 /// Clear every cache owned by this AD context.
145 ///
146 /// # Examples
147 ///
148 /// ```rust
149 /// use tenferro_ad::AdContext;
150 ///
151 /// let ad = AdContext::builder().build().unwrap();
152 /// ad.clear_caches().unwrap();
153 /// assert_eq!(ad.cache_stats().unwrap().ad_transforms.entries, 0);
154 /// ```
155 pub fn clear_caches(&self) -> Result<()> {
156 self.clear_ad_transform_caches()
157 }
158
159 /// Return aggregate cache-entry and retained-byte stats for this AD context.
160 ///
161 /// # Examples
162 ///
163 /// ```rust
164 /// use tenferro_ad::AdContext;
165 ///
166 /// let ad = AdContext::builder().build().unwrap();
167 /// assert_eq!(ad.cache_stats().unwrap().ad_transforms.retained_bytes, 0);
168 /// ```
169 pub fn cache_stats(&self) -> Result<AdContextCacheStats> {
170 Ok(AdContextCacheStats {
171 ad_transforms: self.ad_transform_cache_stats()?,
172 })
173 }
174
175 /// Gradient of a scalar traced output with respect to a traced input.
176 ///
177 /// For complex scalar outputs, tenferro returns the Hermitian-adjoint
178 /// cotangent. To compare seed-`1` scalar gradients with JAX's public
179 /// `grad` values, use the complex conjugate of this result. See
180 /// <https://tensor4all.org/tenferro-rs/guides/complex-ad.html>.
181 ///
182 /// # Examples
183 ///
184 /// ```rust
185 /// use tenferro_ad::AdContext;
186 /// use tenferro_runtime::TracedTensor;
187 ///
188 /// let ad = AdContext::builder().build().unwrap();
189 /// let x = TracedTensor::from_vec_col_major(vec![], vec![3.0_f64]).unwrap();
190 /// let loss = (&x * &x).unwrap();
191 /// let grad = ad.grad(&loss, &x).unwrap();
192 /// assert_eq!(grad.rank, 0);
193 /// ```
194 pub fn grad(&self, output: &TracedTensor, wrt: &TracedTensor) -> Result<TracedTensor> {
195 crate::traced::grad_with_rules_and_cache(
196 output,
197 wrt,
198 &self.extension_rules,
199 Some(self.ad_transform_cache.as_ref()),
200 )
201 }
202
203 /// Gradient that returns `None` when `wrt` is inactive.
204 ///
205 /// # Examples
206 ///
207 /// ```rust
208 /// use tenferro_ad::AdContext;
209 /// use tenferro_runtime::TracedTensor;
210 ///
211 /// let ad = AdContext::builder().build().unwrap();
212 /// let x = TracedTensor::from_vec_col_major(vec![], vec![3.0_f64]).unwrap();
213 /// let loss = (&x * &x).unwrap();
214 /// assert!(ad.grad_optional(&loss, &x).unwrap().is_some());
215 /// ```
216 pub fn grad_optional(
217 &self,
218 output: &TracedTensor,
219 wrt: &TracedTensor,
220 ) -> Result<Option<TracedTensor>> {
221 crate::traced::grad_optional_with_rules_and_cache(
222 output,
223 wrt,
224 &self.extension_rules,
225 Some(self.ad_transform_cache.as_ref()),
226 )
227 }
228
229 /// Forward-mode Jacobian-vector product.
230 ///
231 /// # Examples
232 ///
233 /// ```rust
234 /// use tenferro_ad::AdContext;
235 /// use tenferro_runtime::TracedTensor;
236 ///
237 /// let ad = AdContext::builder().build().unwrap();
238 /// let x = TracedTensor::from_vec_col_major(vec![], vec![3.0_f64]).unwrap();
239 /// let dx = TracedTensor::from_vec_col_major(vec![], vec![1.0_f64]).unwrap();
240 /// let y = (&x * &x).unwrap();
241 /// let dy = ad.jvp(&y, &x, &dx).unwrap();
242 /// assert_eq!(dy.rank, 0);
243 /// ```
244 pub fn jvp(
245 &self,
246 output: &TracedTensor,
247 wrt: &TracedTensor,
248 tangent: &TracedTensor,
249 ) -> Result<TracedTensor> {
250 crate::traced::jvp_with_rules_and_cache(
251 output,
252 wrt,
253 tangent,
254 &self.extension_rules,
255 Some(self.ad_transform_cache.as_ref()),
256 )
257 }
258
259 /// Forward-mode Jacobian-vector product that returns `None` for inactive output.
260 ///
261 /// # Examples
262 ///
263 /// ```rust
264 /// use tenferro_ad::AdContext;
265 /// use tenferro_runtime::TracedTensor;
266 ///
267 /// let ad = AdContext::builder().build().unwrap();
268 /// let x = TracedTensor::from_vec_col_major(vec![], vec![3.0_f64]).unwrap();
269 /// let dx = TracedTensor::from_vec_col_major(vec![], vec![1.0_f64]).unwrap();
270 /// let y = (&x * &x).unwrap();
271 /// assert!(ad.jvp_optional(&y, &x, &dx).unwrap().is_some());
272 /// ```
273 pub fn jvp_optional(
274 &self,
275 output: &TracedTensor,
276 wrt: &TracedTensor,
277 tangent: &TracedTensor,
278 ) -> Result<Option<TracedTensor>> {
279 crate::traced::jvp_optional_with_rules_and_cache(
280 output,
281 wrt,
282 tangent,
283 &self.extension_rules,
284 Some(self.ad_transform_cache.as_ref()),
285 )
286 }
287
288 /// Reverse-mode vector-Jacobian product.
289 ///
290 /// Complex cotangents use tenferro's Hermitian real-inner-product
291 /// convention. Non-real complex cotangent seeds therefore need an explicit
292 /// seed-convention comparison when matching JAX. See
293 /// <https://tensor4all.org/tenferro-rs/guides/complex-ad.html>.
294 ///
295 /// # Examples
296 ///
297 /// ```rust
298 /// use tenferro_ad::AdContext;
299 /// use tenferro_runtime::TracedTensor;
300 ///
301 /// let ad = AdContext::builder().build().unwrap();
302 /// let x = TracedTensor::from_vec_col_major(vec![], vec![3.0_f64]).unwrap();
303 /// let dy = TracedTensor::from_vec_col_major(vec![], vec![1.0_f64]).unwrap();
304 /// let y = (&x * &x).unwrap();
305 /// let dx = ad.vjp(&y, &x, &dy).unwrap();
306 /// assert_eq!(dx.rank, 0);
307 /// ```
308 pub fn vjp(
309 &self,
310 output: &TracedTensor,
311 wrt: &TracedTensor,
312 cotangent: &TracedTensor,
313 ) -> Result<TracedTensor> {
314 crate::traced::vjp_with_rules_and_cache(
315 output,
316 wrt,
317 cotangent,
318 &self.extension_rules,
319 Some(self.ad_transform_cache.as_ref()),
320 )
321 }
322
323 /// Reverse-mode vector-Jacobian product that returns `None` for inactive input.
324 ///
325 /// # Examples
326 ///
327 /// ```rust
328 /// use tenferro_ad::AdContext;
329 /// use tenferro_runtime::TracedTensor;
330 ///
331 /// let ad = AdContext::builder().build().unwrap();
332 /// let x = TracedTensor::from_vec_col_major(vec![], vec![3.0_f64]).unwrap();
333 /// let dy = TracedTensor::from_vec_col_major(vec![], vec![1.0_f64]).unwrap();
334 /// let y = (&x * &x).unwrap();
335 /// assert!(ad.vjp_optional(&y, &x, &dy).unwrap().is_some());
336 /// ```
337 pub fn vjp_optional(
338 &self,
339 output: &TracedTensor,
340 wrt: &TracedTensor,
341 cotangent: &TracedTensor,
342 ) -> Result<Option<TracedTensor>> {
343 crate::traced::vjp_optional_with_rules_and_cache(
344 output,
345 wrt,
346 cotangent,
347 &self.extension_rules,
348 Some(self.ad_transform_cache.as_ref()),
349 )
350 }
351}
352
353/// Builder for [`AdContext`].
354///
355/// # Examples
356///
357/// ```rust
358/// use tenferro_ad::AdContextBuilder;
359///
360/// let ad = AdContextBuilder::new().build().unwrap();
361/// assert!(ad.extension_rules().lookup_linearize("example.missing.v1").is_none());
362/// ```
363#[derive(Clone, Debug, Default)]
364pub struct AdContextBuilder {
365 extension_rule_sets: Vec<ExtensionRuleSet>,
366}
367
368impl AdContextBuilder {
369 /// Create an empty builder.
370 ///
371 /// # Examples
372 ///
373 /// ```rust
374 /// use tenferro_ad::AdContextBuilder;
375 ///
376 /// let _builder = AdContextBuilder::new();
377 /// ```
378 pub fn new() -> Self {
379 Self::default()
380 }
381
382 /// Include an owned extension rule set.
383 ///
384 /// # Examples
385 ///
386 /// ```rust
387 /// use tenferro_ad::{AdContext, extension::ExtensionRuleSet};
388 ///
389 /// let _ad = AdContext::builder()
390 /// .with_extension_rules(ExtensionRuleSet::new())
391 /// .build()
392 /// .unwrap();
393 /// ```
394 pub fn with_extension_rules(mut self, rules: ExtensionRuleSet) -> Self {
395 self.extension_rule_sets.push(rules);
396 self
397 }
398
399 /// Build the context.
400 ///
401 /// Duplicate extension family registrations are rejected.
402 ///
403 /// # Examples
404 ///
405 /// ```rust
406 /// use tenferro_ad::AdContext;
407 ///
408 /// let ad = AdContext::builder().build().unwrap();
409 /// assert!(ad.extension_rules().lookup_linearize("example.missing.v1").is_none());
410 /// ```
411 pub fn build(self) -> std::result::Result<AdContext, ExtensionRegistryError> {
412 let mut extension_rules = ExtensionRuleSet::new();
413 for rules in self.extension_rule_sets {
414 extension_rules.merge(rules)?;
415 }
416 Ok(AdContext {
417 extension_rules,
418 ad_transform_cache: Arc::new(AdTransformCache::new()),
419 })
420 }
421}