Skip to main content

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`s are generic over the backend to allow users to perform operations using different `Backend` implementations.
7//! Burn's tensors also support auto-differentiation thanks to the `AutodiffBackend` trait.
8//!
9//! # Note for contributors: `*_impl` helpers
10//!
11//! Throughout this crate (e.g. in `tensor::api::float`, `tensor::api::int`,
12//! `tensor::api::bool`, `tensor::api::cast`, and `tensor::activation`), public
13//! generic methods on `Tensor<D, K>` that need to call into `burn_dispatch` are
14//! routed through small non-generic helper functions named `*_impl`, grouped
15//! together at the bottom of each file under a banner like:
16//!
17//! ```text
18//! // =====================================================================
19//! // Non-generic implementation helpers (outlined from the generic API).
20//! // =====================================================================
21//! ```
22//!
23//! These helpers take and return only `BridgeTensor` (a type-erased blob — no
24//! `DispatchTensor` or other `burn_dispatch` types appear in their
25//! signatures). Because the helpers are not generic over `D`, they are
26//! compiled once, and the MIR of the public generic methods does not mention
27//! any `burn_dispatch` types. Downstream crates that monomorphize the public
28//! generic API therefore never have to resolve the cubecl-backed type tree,
29//! which drastically cuts compile times for user code.
30//!
31//! When adding a new public method that calls a `Dispatch::*` op, follow the
32//! same pattern: keep the generic method body thin and forward to a
33//! non-generic `*_impl` helper alongside the existing ones.
34
35#[macro_use]
36extern crate derive_new;
37
38extern crate alloc;
39
40mod bridge;
41mod tensor;
42
43pub(crate) use tensor::check::macros::check;
44pub use tensor::*;
45
46// Re-exported types
47pub use burn_std::{
48    AllocationProperty, Bytes, bf16, f16,
49    reader::{read_sync, try_read_sync},
50    stream_id::StreamId,
51};
52
53mod device;
54pub use device::*;
55
56#[cfg(feature = "remote-server")]
57pub mod server;
58
59pub(crate) use burn_backend::TensorPrimitive;