burn-tensor 0.22.0-pre.4

Tensor library with user-friendly APIs and automatic differentiation support
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68

#![cfg_attr(not(feature = "std"), no_std)]
#![warn(missing_docs)]
#![cfg_attr(docsrs, feature(doc_cfg))]

//! This library provides the core abstractions required to run tensor operations with Burn.
//! [`Tensor<D, K>`](Tensor) is generic over rank and tensor kind; its [`Device`] selects the
//! backend at runtime. Operations pass through an opaque bridge and runtime dispatch to backend
//! primitives. With the `autodiff` feature, create tensors on `Device::autodiff()` devices and
//! mark source leaves with [`Tensor::require_grad`] to record gradients.
//!
//! # Note for contributors: `*_impl` helpers
//!
//! Throughout this crate (e.g. in `tensor::api::float`, `tensor::api::int`,
//! `tensor::api::bool`, `tensor::api::cast`, and `tensor::activation`), public
//! generic methods on `Tensor<D, K>` that need to call into `burn_dispatch` are
//! routed through small non-generic helper functions named `*_impl`, grouped
//! together at the bottom of each file under a banner like:
//!
//! ```text
//! // =====================================================================
//! // Non-generic implementation helpers (outlined from the generic API).
//! // =====================================================================
//! ```
//!
//! These helpers take and return only `BridgeTensor` (a type-erased blob — no
//! `DispatchTensor` or other `burn_dispatch` types appear in their
//! signatures). Because the helpers are not generic over `D`, they are
//! compiled once, and the MIR of the public generic methods does not mention
//! any `burn_dispatch` types. Downstream crates that monomorphize the public
//! generic API therefore never have to resolve the cubecl-backed type tree,
//! which drastically cuts compile times for user code.
//!
//! When adding a new public method that calls a `Dispatch::*` op, follow the
//! same pattern: keep the generic method body thin and forward to a
//! non-generic `*_impl` helper alongside the existing ones.

#[macro_use]
extern crate derive_new;

extern crate alloc;

mod bridge;
mod tensor;

pub(crate) use tensor::check::macros::check;
pub use tensor::*;

mod einsum_macros;
mod shape_macros;
#[doc(hidden)]
pub use burn_derive::{__assert_shape, __debug_assert_shape, __einsum};

// Re-exported types
#[cfg(feature = "autodiff")]
pub use burn_dispatch::GradientCheckpointingStrategy;
pub use burn_std::{
    AllocationProperty, Bytes, bf16, f16,
    reader::{read_sync, try_read_sync},
    stream::StreamId,
};

mod device;
pub use device::*;

#[cfg(feature = "remote-server")]
pub mod server;

pub(crate) use burn_backend::TensorPrimitive;