burn_tensor/lib.rs
1#![cfg_attr(not(feature = "std"), no_std)]
2#![warn(missing_docs)]
3#![cfg_attr(docsrs, feature(doc_cfg))]
4
5//! This library provides the core abstractions required to run tensor operations with Burn.
6//! [`Tensor<D, K>`](Tensor) is generic over rank and tensor kind; its [`Device`] selects the
7//! backend at runtime. Operations pass through an opaque bridge and runtime dispatch to backend
8//! primitives. With the `autodiff` feature, create tensors on `Device::autodiff()` devices and
9//! mark source leaves with [`Tensor::require_grad`] to record gradients.
10//!
11//! # Note for contributors: `*_impl` helpers
12//!
13//! Throughout this crate (e.g. in `tensor::api::float`, `tensor::api::int`,
14//! `tensor::api::bool`, `tensor::api::cast`, and `tensor::activation`), public
15//! generic methods on `Tensor<D, K>` that need to call into `burn_dispatch` are
16//! routed through small non-generic helper functions named `*_impl`, grouped
17//! together at the bottom of each file under a banner like:
18//!
19//! ```text
20//! // =====================================================================
21//! // Non-generic implementation helpers (outlined from the generic API).
22//! // =====================================================================
23//! ```
24//!
25//! These helpers take and return only `BridgeTensor` (a type-erased blob — no
26//! `DispatchTensor` or other `burn_dispatch` types appear in their
27//! signatures). Because the helpers are not generic over `D`, they are
28//! compiled once, and the MIR of the public generic methods does not mention
29//! any `burn_dispatch` types. Downstream crates that monomorphize the public
30//! generic API therefore never have to resolve the cubecl-backed type tree,
31//! which drastically cuts compile times for user code.
32//!
33//! When adding a new public method that calls a `Dispatch::*` op, follow the
34//! same pattern: keep the generic method body thin and forward to a
35//! non-generic `*_impl` helper alongside the existing ones.
36
37#[macro_use]
38extern crate derive_new;
39
40extern crate alloc;
41
42mod bridge;
43mod tensor;
44
45pub(crate) use tensor::check::macros::check;
46pub use tensor::*;
47
48mod einsum_macros;
49mod shape_macros;
50#[doc(hidden)]
51pub use burn_derive::{__assert_shape, __debug_assert_shape, __einsum};
52
53// Re-exported types
54#[cfg(feature = "autodiff")]
55pub use burn_dispatch::GradientCheckpointingStrategy;
56pub use burn_std::{
57 AllocationProperty, Bytes, bf16, f16,
58 reader::{read_sync, try_read_sync},
59 stream::StreamId,
60};
61
62mod device;
63pub use device::*;
64
65#[cfg(feature = "remote-server")]
66pub mod server;
67
68pub(crate) use burn_backend::TensorPrimitive;