Skip to main content

oxmera_tensor/
lib.rs

1//! The oxmera tensor: shared storage viewed through a layout, with
2//! multi-device backends and reverse-mode autograd.
3//!
4//! This crate is the hub of the framework:
5//!
6//! - [`Tensor`] — the value type: constructors, zero-copy strided views,
7//!   element access, and the full differentiable op surface (`add`,
8//!   `matmul`, `softmax`, …) plus `std::ops` operator sugar.
9//! - [`einsum`] — Einstein-summation contractions lowered onto `matmul`,
10//!   `permute` and `sum`.
11//! - Small batched linear algebra on `Tensor`: `eye`, `diag`, `diag_embed`,
12//!   `trace`, `cholesky`, `logdet`, `det`, `eigh`.
13//! - [`backend`] — the op vocabulary ([`backend::UnaryOp`],
14//!   [`backend::BinaryOp`], [`backend::ReduceOp`]), the [`backend::Backend`]
15//!   trait every device implements, and the registry that resolves a
16//!   [`oxmera_core::Device`] handle. The traits live here so tensor
17//!   methods and operator overloads can dispatch without violating the
18//!   orphan rule.
19//! - [`autograd`] — the tape: recording switch, [`autograd::no_grad`],
20//!   and gradient propagation; `Tensor::backward` drives it.
21//!
22//! Backends register themselves at load time (linking `oxmera-cpu` or
23//! `oxmera-metal` is what makes their device usable); the `oxmera`
24//! umbrella crate links every backend for the current platform.
25
26#![deny(unsafe_code)] // the two Metal Send/Sync impls in `storage` opt in locally
27#![warn(missing_docs)]
28
29pub mod autograd;
30pub mod backend;
31pub mod cpu;
32mod cpu_f64;
33mod cpu_iter;
34mod cpu_linalg;
35mod cpu_matmul;
36mod einsum;
37mod linalg;
38pub use linalg::EIGH_SYMMETRY_TOL;
39pub mod ops;
40pub mod overload;
41pub mod storage;
42pub mod tensor;
43
44pub use autograd::{NoGradGuard, no_grad};
45pub use backend::{Backend, BinaryOp, ReduceOp, UnaryOp, backend_for, register_backend};
46pub use einsum::einsum;
47pub use storage::{CpuStorage, OpaqueBuffer, Storage, StorageData};
48pub use tensor::Tensor;
49
50/// What this crate can do, for `oxmera doctor`.
51///
52/// It lives here, beside the code, because the alternative was a list of
53/// string literals in `oxmera-cli` that nobody editing an op ever opened:
54/// that list named neither CUDA nor `f64` nor the linear algebra two
55/// releases after they shipped, and three golden files pinned it, so
56/// adding a feature to it cost more than leaving it stale.
57///
58/// Adding an op family here is one line in the crate that implements it.
59pub const CAPABILITIES: &[(&str, &str)] = &[
60    // Each row must fit an 80-column terminal after the 12-character
61    // "  area     " prefix, so keep the text under 68 characters: the
62    // doctor goldens are 100 wide and a wrapped row costs two of them.
63    // `capability_rows_fit_a_narrow_terminal` in oxmera-cli enforces it.
64    (
65        "tensor",
66        "strided views, broadcasting, batched matmul, device transfer",
67    ),
68    (
69        "dtypes",
70        "f32 everywhere; f64 on the CPU; i64 indices; no promotion",
71    ),
72    (
73        "linalg",
74        "eye diag trace, batched cholesky (differentiable) logdet det eigh",
75    ),
76    (
77        "einsum",
78        "one- and two-operand contractions with an explicit output",
79    ),
80    (
81        "indexing",
82        "index_select / index_add, native on Metal and CUDA",
83    ),
84];