burn_core/
lib.rs

1#![cfg_attr(not(feature = "std"), no_std)]
2#![warn(missing_docs)]
3#![cfg_attr(docsrs, feature(doc_cfg))]
4
5//! The core of Burn: modules, configuration, records and data loading, on top of the tensor API.
6//!
7//! Applications should depend on [`burn`](https://docs.rs/burn), which re-exports everything here
8//! together with neural network layers, optimizers and training. This crate is the dependency for
9//! libraries that only need the core abstractions.
10//!
11//! - [`tensor`]: Burn's tensor API, re-exported from `burn-tensor`.
12//! - [`module`]: the [`Module`](module::Module) trait, implemented with `#[derive(Module)]`, and
13//!   [`Param`](module::Param) for trainable tensors.
14//! - [`config`]: serializable configuration structs with `#[derive(Config)]`.
15//! - [`store`]: module records and the burnpack format.
16//! - [`data`]: datasets, batchers and data loaders (`std` only).
17//! - [`prelude`]: the types most programs import.
18//!
19//! # Feature flags
20//!
21//! Backend features (`wgpu`, `cuda`, `flex`, ...) and the `autodiff`, `fusion`, `remote` and
22//! `capture` features match those of `burn`. Others:
23//!
24//! - `std` (default): standard library support. Without it the crate is `no_std` with `alloc`.
25//! - `dataset`: the dataset library in [`data`]; `vision`, `audio` and `sqlite` add sources.
26//! - `network`: file downloads with a progress bar.
27//! - `tracing`: instrument operations with the `tracing` crate.
28
29// `derive_new` provides the `#[derive(new)]` macro used across the crate; the lint mistakenly
30// reports the `#[macro_use]` as unused.
31#[allow(unused_imports)]
32#[macro_use]
33extern crate derive_new;
34
35/// Re-export serde for proc macros.
36pub use serde;
37
38/// The configuration module.
39pub mod config;
40
41/// Data module.
42#[cfg(feature = "std")]
43pub mod data;
44
45/// Module for the neural network module.
46pub mod module;
47
48/// Module for saving and loading module/optimizer state in the burnpack format.
49pub mod store;
50
51/// Module for the tensor.
52pub mod tensor;
53// Tensor at root: `burn::Tensor`
54pub use tensor::Tensor;
55
56#[cfg(feature = "extension")]
57/// Backend module.
58pub mod backend;
59
60extern crate alloc;
61
62// TODO: configurable device priority
63#[cfg(test)]
64#[allow(missing_docs)]
65pub fn test_device() -> burn_tensor::Device {
66    burn_tensor::Device::flex()
67}
68
69#[cfg(test)]
70mod test_utils {
71    use crate as burn;
72    use crate::module::Module;
73    use crate::module::Param;
74    use burn_tensor::Device;
75    use burn_tensor::Tensor;
76
77    /// Simple linear module.
78    #[derive(Module, Debug)]
79    pub struct SimpleLinear {
80        pub weight: Param<Tensor<2>>,
81        pub bias: Option<Param<Tensor<1>>>,
82    }
83
84    impl SimpleLinear {
85        pub fn new(in_features: usize, out_features: usize, device: &Device) -> Self {
86            let weight = Tensor::random(
87                [out_features, in_features],
88                burn_tensor::Distribution::Default,
89                device,
90            );
91            let bias = Tensor::random([out_features], burn_tensor::Distribution::Default, device);
92
93            Self {
94                weight: Param::from_tensor(weight),
95                bias: Some(Param::from_tensor(bias)),
96            }
97        }
98    }
99}
100
101pub mod prelude {
102    //! Structs and macros used by most projects. Add `use
103    //! burn::prelude::*` to your code to quickly get started with
104    //! Burn.
105    pub use crate::{
106        config::Config,
107        module::Module,
108        tensor::{
109            Bool, Device, DeviceIndex, DeviceKind, ElementConversion, Float, Int, Shape, SliceArg,
110            Tensor, TensorData, assert_shape, cast::ToElement, debug_assert_shape, s,
111        },
112    };
113}
burn_core/lib.rs

burn_core/
lib.rs
