tenferro-ad 0.2.0

Eager runtime, eager tensors, and traced AD extension traits for tenferro.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71

//! Automatic differentiation APIs for tenferro.
//!
//! This crate is the explicit opt-in boundary for traced and eager automatic
//! differentiation. Primal graph construction and execution live in
//! `tenferro-runtime`; tensor storage lives in `tenferro-tensor`, and CPU
//! execution lives in `tenferro-cpu`.
//!
//! Use [`EagerRuntime`] and [`EagerTensor`] for PyTorch-style immediate
//! execution where tracked variables accumulate gradients after `backward()`.
//! Use [`TracedTensorAdExt`] or [`AdContext`] for JAX-style graph transforms
//! such as `grad`, `vjp`, and `jvp` on [`tenferro_runtime::TracedTensor`]
//! values. `AdContext` is the explicit place to add extension AD rule sets for
//! operation-family crates such as `tenferro-linalg`.
//!
//! User-facing guides live at
//! <https://tensor4all.org/tenferro-rs/guides/autodiff.html> and
//! <https://tensor4all.org/tenferro-rs/guides/choosing-an-api.html>.
//!
//! # Examples
//!
//! ```rust
//! use tenferro_ad::AdContext;
//! use tenferro_runtime::TracedTensor;
//!
//! let ad = AdContext::builder().build().unwrap();
//! let x = TracedTensor::from_vec_col_major(vec![], vec![3.0_f64]).unwrap();
//! let loss = (&x * &x).unwrap();
//! let dx = ad.grad(&loss, &x).unwrap();
//! assert_eq!(dx.rank, 0);
//! ```

mod ad_rule_error;
mod context;
mod eager;
mod eager_backend;
mod eager_builder;
pub(crate) mod eager_exec;
pub(crate) mod eager_ops;
pub(crate) mod eager_ops_elementwise;
pub mod extension;
mod shape_packing;
pub mod traced;

pub use context::{AdContext, AdContextBuilder};
pub use eager::{EagerRuntime, EagerRuntimeCacheStats, EagerTensor};
pub use eager_backend::EagerBackend;
pub(crate) use tenferro_runtime::{extension_cache, extension_runtime, scalar_semantics};
pub(crate) mod shape_infer {
    pub use tenferro_runtime::extension::{
        promote_dtype, promote_dtype_for_binary_op, promote_dtypes,
    };
}
pub use tenferro_runtime::{
    CompareDir, DType, DotGeneralConfig, GatherConfig, PadConfig, ScatterConfig, SliceConfig,
    Tensor,
};
pub use traced::TracedTensorAdExt;

pub use tenferro_runtime::{ContextId, Error, Result};

pub mod error {
    pub use tenferro_runtime::{ContextId, Error, Result};
}

pub(crate) mod metadata {
    pub use tenferro_runtime::ad_support::{
        metadata_scopes_for_scope, push_metadata_scope, register_scoped_live_graph_metadata,
        register_scoped_metadata_batch, register_scoped_value_metadata, tensor_meta_from_tensor,
        GlobalMetadataScope,
    };
}