tenferro_ad/
lib.rs

1//! Automatic differentiation APIs for tenferro.
2//!
3//! This crate is the explicit opt-in boundary for traced and eager automatic
4//! differentiation. Primal graph construction and execution live in
5//! `tenferro-runtime`; tensor storage lives in `tenferro-tensor`, and CPU
6//! execution lives in `tenferro-cpu`.
7//!
8//! Use [`EagerRuntime`] and [`EagerTensor`] for PyTorch-style immediate
9//! execution where tracked variables accumulate gradients after `backward()`.
10//! Use [`TracedTensorAdExt`] or [`AdContext`] for JAX-style graph transforms
11//! such as `grad`, `vjp`, and `jvp` on [`tenferro_runtime::TracedTensor`]
12//! values. `AdContext` is the explicit place to add extension AD rule sets for
13//! operation-family crates such as `tenferro-linalg`.
14//!
15//! User-facing guides live at
16//! <https://tensor4all.org/tenferro-rs/guides/autodiff.html> and
17//! <https://tensor4all.org/tenferro-rs/guides/choosing-an-api.html>.
18//!
19//! # Examples
20//!
21//! ```rust
22//! use tenferro_ad::AdContext;
23//! use tenferro_runtime::TracedTensor;
24//!
25//! let ad = AdContext::builder().build().unwrap();
26//! let x = TracedTensor::from_vec_col_major(vec![], vec![3.0_f64]).unwrap();
27//! let loss = (&x * &x).unwrap();
28//! let dx = ad.grad(&loss, &x).unwrap();
29//! assert_eq!(dx.rank, 0);
30//! ```
31
32mod context;
33mod eager;
34mod eager_backend;
35pub(crate) mod eager_exec;
36pub(crate) mod eager_ops;
37pub(crate) mod eager_ops_elementwise;
38pub mod extension;
39// semantic_compat removed in Unification 7.
40pub mod semantic_extension;
41pub mod semantic_transform;
42mod shape_packing;
43pub mod traced;
44mod transform_cache;
45
46pub use context::{AdContext, AdContextBuilder, AdContextCacheStats};
47pub use eager::{
48    CpuPlacementBoundEager, EagerNoGradGuard, EagerRuntime, EagerRuntimeCacheStats, EagerTensor,
49};
50pub use shape_packing::EagerSliceBuilder;
51pub(crate) use tenferro_runtime::{extension_cache, scalar_semantics};
52pub use transform_cache::AdTransformCacheLimits;
53pub(crate) mod shape_infer {
54    pub use tenferro_runtime::extension::{
55        promote_dtype, promote_dtype_for_binary_op, promote_dtypes,
56    };
57}
58pub use tenferro_runtime::{
59    CompareDir, DType, DotGeneralConfig, GatherConfig, PadConfig, ScatterConfig, SliceConfig,
60    Tensor,
61};
62pub use traced::TracedTensorAdExt;
63
64pub use tenferro_runtime::{ContextId, Error, Result};
65
66pub mod error {
67    pub use tenferro_runtime::{ContextId, Error, Result};
68}
69
70pub(crate) mod metadata {
71    pub use tenferro_runtime::ad_support::{
72        metadata_scopes_for_scope, push_metadata_scope, register_scoped_metadata_batch,
73        register_scoped_value_metadata, tensor_meta_from_tensor, GlobalMetadataScope,
74    };
75}
tenferro_ad/lib.rs

tenferro_ad/
lib.rs
