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//!
7//! # Examples
8//!
9//! ```
10//! use tenferro_linalg::TracedTensorLinalgExt;
11//! use tenferro_cpu::CpuEngine;
12//! use tenferro_runtime::{ExecutionEngine, GraphCompiler, Runtime, TracedTensor};
13//!
14//! let a = TracedTensor::from_vec_col_major(
15//!     vec![2, 2],
16//!     vec![4.0_f64, 2.0, 2.0, 3.0],
17//! )
18//! .unwrap();
19//! let l = a.cholesky().unwrap();
20//!
21//! let mut compiler = GraphCompiler::new();
22//! let program = compiler.compile(&l).unwrap();
23//! let mut backend = tenferro_cpu::CpuEngine::serial();
24//! let engine_id = tenferro_cpu::runtime_engine_id().unwrap();
25//! let mut builder = Runtime::builder();
26//! builder
27//!     .register_engine(tenferro_cpu::runtime_engine_registration(&backend).unwrap())
28//!     .unwrap();
29//! builder
30//!     .install_extension_module(tenferro_linalg::extension_module::<CpuEngine>(engine_id).unwrap())
31//!     .unwrap();
32//! let runtime = builder.build().unwrap();
33//! let prepared = runtime.prepare(&program, &[], &[backend.target()]).unwrap();
34//! let out = {
35//!     let mut engines: [&mut dyn ExecutionEngine; 1] = [&mut backend];
36//!     runtime.run(&mut engines, &prepared, &[]).unwrap().pop().unwrap()
37//! }
38//! .into_tensor(&mut backend)
39//! .unwrap();
40//! assert_eq!(out.shape(), &[2, 2]);
41//! ```
42//!
43//! # Cargo features
44//!
45//! | Feature | Enables |
46//! |---|---|
47//! | `native` (default) and the other CPU provider features | Forwarded to `tenferro-cpu`; see its documentation. |
48//! | `autodiff` | The eager surface (`EagerSessionLinalgExt`, `EagerTensorLinalgExt`) and AD rules. Adds the `tenferro-ad` dependency. |
49//! | `cuda` | CUDA execution through `tenferro-gpu`. |
50//! | `rocm` | Placeholder; HIP/ROCm is not implemented. |
51//!
52//! `autodiff` is not a heavy dependency switch: `tenferro-ad` is the crate that
53//! owns `EagerSession`/`EagerTensor`, so an eager linalg surface without it
54//! would save nothing and would only produce eager tensors whose linalg ops
55//! fail on backward (#1972). For inference without AD, call the same
56//! operations on concrete tensors inside a backend session through
57//! [`TensorLinalgExt`] / [`TypedTensorLinalgExt`], which need no `autodiff`:
58//!
59//! ```
60//! use tenferro_cpu::CpuEngine;
61//! use tenferro_linalg::TypedTensorLinalgExt;
62//! use tenferro_tensor::TypedTensor;
63//!
64//! let mut backend = tenferro_cpu::CpuEngine::serial();
65//! // Lower-triangular [[2, 0], [1, 1]] (column-major) and right-hand side.
66//! let l = TypedTensor::<f64>::from_vec_col_major(vec![2, 2], vec![2.0, 1.0, 0.0, 1.0])?;
67//! let b = TypedTensor::<f64>::from_vec_col_major(vec![2, 1], vec![4.0, 5.0])?;
68//! let x = {
69//!     // left_side, lower, transpose_a, unit_diagonal
70//!     l.triangular_solve(&b, true, true, false, false, &mut backend)
71//! }?;
72//! assert_eq!(x.host_data()?, &[2.0, 3.0]);
73//! # Ok::<(), Box<dyn std::error::Error>>(())
74//! ```
75// A misaligned pointer handed to a CUDA library is undefined behaviour and
76// fails only on some library versions: `cuDoubleComplex` is `double2`, which
77// CUDA declares `__align__(16)`, while `num_complex::Complex64` is 8-aligned,
78// so casting `&Complex64` to `*const cuDoubleComplex` produced a pointer
79// cuBLAS >= 12.9 faults on (issue #1870). Deny the whole cast class at this
80// FFI boundary rather than re-auditing it by hand.
81#![cfg_attr(docsrs, feature(doc_cfg))]
82#![deny(clippy::cast_ptr_alignment)]
83
84#[cfg(feature = "autodiff")]
85mod ad;
86pub mod backend;
87mod cpu;
88#[cfg(feature = "autodiff")]
89mod eager_composites;
90#[cfg(feature = "autodiff")]
91mod eager_ext;
92pub mod error;
93mod extension;
94#[cfg(feature = "cuda")]
95mod gpu;
96mod householder;
97pub mod prelude;
98mod rank_revealing_qr;
99mod tensor_ext;
100mod traced;
101mod validation;
102
103#[cfg(feature = "autodiff")]
104pub use ad::semantic_ad_rules;
105#[cfg(feature = "autodiff")]
106pub use ad::support::{
107    all_linalg_ad_support, linalg_ad_support, LinalgAdModeSupport, LinalgAdOpKind,
108    LinalgAdOutputSupport, LinalgAdRoute, LinalgAdRuleSupport, LinalgAdSupport,
109};
110pub use backend::LinalgBackend;
111#[cfg(feature = "autodiff")]
112#[cfg_attr(docsrs, doc(cfg(feature = "autodiff")))]
113pub use eager_ext::{EagerSessionLinalgExt, EagerTensorLinalgExt};
114// The CPU linalg route is compiled per backend; a build with neither is a
115// configuration error, named as such rather than as a missing internal type.
116#[cfg(not(any(feature = "native", feature = "blas")))]
117compile_error!(
118    "tenferro-linalg needs exactly one CPU backend: enable the `native` (default) \
119     or `blas` feature"
120);
121
122pub use error::{Error, Result};
123pub use extension::{
124    extension_module, EighDriver, EighGauge, EighOptions, QrGauge, QrOptions, SvdDriver, SvdGauge,
125    SvdOptions, DEFAULT_DECOMPOSITION_DERIVATIVE_EPS, LINALG_EXTENSION_FAMILY_ID,
126};
127pub use householder::HouseholderQr;
128pub use rank_revealing_qr::{RankRevealingQrOptions, RankRevealingQrResult};
129pub use tensor_ext::{
130    LinalgScalar, TensorLinalgExt, TensorReadLinalgExt, TypedEig, TypedFullPivLu, TypedLu,
131    TypedRankRevealingQrResult, TypedSvd, TypedTensorLinalgExt,
132};
133pub use traced::TracedTensorLinalgExt;
tenferro_linalg/lib.rs

tenferro_linalg/
lib.rs
