burn_backend/
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 types that define how Burn tensor data is represented, stored, and interpreted.
6
7#[macro_use]
8extern crate derive_new;
9
10extern crate alloc;
11
12mod data;
13pub use data::*;
14
15pub mod distribution;
16pub use distribution::*;
17pub mod element;
18pub use element::*;
19
20/// [`Backend`] trait and required types.
21pub mod backend;
22pub use backend::*;
23
24/// Backend tensor primitives and operations.
25pub mod tensor;
26
27// Re-exported types
28pub use burn_std::reader::*; // Useful so that backends don't have to add `burn_std` as a dependency.
29pub use burn_std::{
30    AllocationProperty, BoolDType, BoolStore, Bytes, DType, DeviceHandle, FloatDType, IntDType,
31    bf16, f16, stream_id::StreamId,
32};
33
34/// Shape definition.
35pub mod shape {
36    pub use burn_std::shape::*;
37}
38pub use shape::*;
39
40/// Slice utilities.
41pub mod slice {
42    pub use burn_std::{s, slice::*};
43}
44pub use slice::*;
45
46/// Indexing utilities.
47pub mod indexing {
48    pub use burn_std::indexing::*;
49}
50pub use indexing::*;
51
52/// Quantization data representation.
53pub mod quantization {
54    pub use crate::tensor::quantization::*;
55    pub use burn_std::quantization::{
56        BlockSize, QuantLevel, QuantMode, QuantParam, QuantPropagation, QuantScheme, QuantStore,
57        QuantValue, QuantizedBytes,
58    };
59}
60
61#[cfg(feature = "cubecl-wgpu")]
62mod cube_wgpu {
63    use crate::backend::DeviceOps;
64    use cubecl::wgpu::WgpuDevice;
65
66    impl DeviceOps for WgpuDevice {}
67}
68
69#[cfg(feature = "cubecl-cuda")]
70mod cube_cuda {
71    use crate::backend::DeviceOps;
72    use cubecl::cuda::CudaDevice;
73
74    impl DeviceOps for CudaDevice {}
75}
76
77#[cfg(feature = "cubecl-cpu")]
78mod cube_cpu {
79    use crate::backend::DeviceOps;
80    use cubecl::cpu::CpuDevice;
81
82    impl DeviceOps for CpuDevice {}
83}
84
85#[cfg(feature = "cubecl-hip")]
86mod cube_hip {
87    use crate::backend::DeviceOps;
88    use cubecl::hip::AmdDevice;
89
90    impl DeviceOps for AmdDevice {}
91}
92
93/// Convenience macro to link to the `burn-tensor` docs for this crate version.
94///
95/// Usage:
96/// ```rust,ignore
97/// # use burn_backend::doc_tensor;
98/// doc_tensor!();        // Links to `Tensor` struct
99/// doc_tensor!("zeros"); // Links to `Tensor::zeros` method
100/// ```
101#[macro_export]
102macro_rules! doc_tensor {
103    () => {
104        concat!(
105            "[`Tensor`](https://docs.rs/burn-tensor/",
106            env!("CARGO_PKG_VERSION"),
107            "/burn_tensor/struct.Tensor.html)"
108        )
109    };
110
111    ($method:literal) => {
112        concat!(
113            "[`Tensor::",
114            $method,
115            "`](",
116            "https://docs.rs/burn-tensor/",
117            env!("CARGO_PKG_VERSION"),
118            "/burn_tensor/struct.Tensor.html#method.",
119            $method,
120            ")"
121        )
122    };
123}
burn_backend/lib.rs

burn_backend/
lib.rs
