tenferro_linalg/lib.rs
1//! Linear algebra extension operations for tenferro.
2//!
3//! This crate owns the graph-facing linalg op payloads and runtime
4//! registration. Tensor-facing operations are exposed through extension traits.
5//! CPU backend kernels live in this crate behind the linalg backend trait.
6//! A CPU backend paired by `tenferro_gpu::AppleContext` additionally supports
7//! guarded rank-2 Cholesky on matching Apple managed `F32`, `F64`, `C32`, and
8//! `C64` tensors. This is an explicit CPU selection and is not a general
9//! managed-memory fallback for other linalg operations.
10//!
11//! # Examples
12//!
13//! ```
14//! use tenferro_linalg::TracedTensorLinalgExt;
15//! use tenferro_cpu::CpuBackend;
16//! use tenferro_runtime::{GraphCompiler, Runtime, TracedTensor};
17//!
18//! let a = TracedTensor::from_vec_col_major(
19//! vec![2, 2],
20//! vec![4.0_f64, 2.0, 2.0, 3.0],
21//! )
22//! .unwrap();
23//! let l = a.cholesky().unwrap();
24//!
25//! let mut compiler = GraphCompiler::new();
26//! let program = compiler.compile(&l).unwrap();
27//! let backend = CpuBackend::new();
28//! let engine_id = tenferro_cpu::runtime_engine_id().unwrap();
29//! let mut builder = Runtime::builder();
30//! builder
31//! .register_engine(tenferro_cpu::runtime_engine_registration(&backend).unwrap())
32//! .unwrap();
33//! builder
34//! .install_extension_module(tenferro_linalg::extension_module::<CpuBackend>(engine_id).unwrap())
35//! .unwrap();
36//! let runtime = builder.build().unwrap();
37//! let out = runtime.run_compiled(&program, &[]).unwrap().pop().unwrap();
38//! assert_eq!(out.shape(), &[2, 2]);
39//! ```
40
41#[cfg(feature = "autodiff")]
42mod ad;
43pub mod backend;
44mod cpu;
45#[cfg(feature = "autodiff")]
46mod eager_composites;
47#[cfg(feature = "autodiff")]
48mod eager_ext;
49pub mod error;
50mod extension;
51#[cfg(feature = "cuda")]
52mod gpu;
53mod tensor_ext;
54mod traced;
55
56#[cfg(feature = "autodiff")]
57pub use ad::semantic_ad_rules;
58#[cfg(feature = "autodiff")]
59pub use ad::support::{
60 all_linalg_ad_support, linalg_ad_support, LinalgAdModeSupport, LinalgAdOpKind,
61 LinalgAdOutputSupport, LinalgAdRoute, LinalgAdRuleSupport, LinalgAdSupport,
62};
63pub use backend::LinalgBackend;
64#[cfg(feature = "autodiff")]
65pub use eager_ext::EagerTensorLinalgExt;
66pub use error::{Error, Result};
67pub use extension::{
68 extension_module, EighGauge, EighOptions, QrGauge, QrOptions, SvdGauge, SvdOptions,
69 DEFAULT_DECOMPOSITION_DERIVATIVE_EPS, LINALG_EXTENSION_FAMILY_ID,
70};
71pub use tensor_ext::{
72 LinalgScalar, TensorLinalgExt, TensorReadLinalgExt, TypedEig, TypedFullPivLu, TypedLu,
73 TypedSvd, TypedTensorLinalgExt,
74};
75pub use traced::TracedTensorLinalgExt;