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, Bytes, DType, FloatDType, IntDType, bf16, f16, stream_id::StreamId,
31};
32
33/// Shape definition.
34pub mod shape {
35    pub use burn_std::shape::*;
36}
37pub use shape::*;
38
39/// Slice utilities.
40pub mod slice {
41    pub use burn_std::{s, slice::*};
42}
43pub use slice::*;
44
45/// Indexing utilities.
46pub mod indexing {
47    pub use burn_std::indexing::*;
48}
49pub use indexing::*;
50
51/// Quantization data representation.
52pub mod quantization {
53    pub use crate::tensor::quantization::*;
54    pub use burn_std::quantization::{
55        BlockSize, QuantLevel, QuantMode, QuantParam, QuantPropagation, QuantScheme, QuantStore,
56        QuantValue, QuantizedBytes,
57    };
58}
59
60#[cfg(feature = "cubecl-wgpu")]
61mod cube_wgpu {
62    use crate::backend::DeviceOps;
63    use cubecl::wgpu::WgpuDevice;
64
65    impl DeviceOps for WgpuDevice {}
66}
67
68#[cfg(feature = "cubecl-cuda")]
69mod cube_cuda {
70    use crate::backend::DeviceOps;
71    use cubecl::cuda::CudaDevice;
72
73    impl DeviceOps for CudaDevice {}
74}
75
76#[cfg(all(feature = "cubecl-cpu", target_os = "linux"))]
77mod cube_cpu {
78    use crate::backend::DeviceOps;
79    use cubecl::cpu::CpuDevice;
80
81    impl DeviceOps for CpuDevice {}
82}
83
84#[cfg(feature = "cubecl-hip")]
85mod cube_hip {
86    use crate::backend::DeviceOps;
87    use cubecl::hip::AmdDevice;
88
89    impl DeviceOps for AmdDevice {}
90}
91
92/// Convenience macro to link to the `burn-tensor` docs for this crate version.
93///
94/// Usage:
95/// ```rust,ignore
96/// # use burn_backend::doc_tensor;
97/// doc_tensor!();        // Links to `Tensor` struct
98/// doc_tensor!("zeros"); // Links to `Tensor::zeros` method
99/// ```
100#[macro_export]
101macro_rules! doc_tensor {
102    () => {
103        concat!(
104            "[`Tensor`](https://docs.rs/burn-tensor/",
105            env!("CARGO_PKG_VERSION"),
106            "/burn_tensor/struct.Tensor.html)"
107        )
108    };
109
110    ($method:literal) => {
111        concat!(
112            "[`Tensor::",
113            $method,
114            "`](",
115            "https://docs.rs/burn-tensor/",
116            env!("CARGO_PKG_VERSION"),
117            "/burn_tensor/struct.Tensor.html#method.",
118            $method,
119            ")"
120        )
121    };
122}
burn_backend/lib.rs

burn_backend/
lib.rs
