burn_derive/
lib.rs

1#![warn(missing_docs)]
2
3//! Derive and procedural macros for Burn.
4//!
5//! Use these through `burn`, which re-exports them; this crate is not meant to be a direct
6//! dependency.
7//!
8//! - `#[derive(Module)]` implements `Module` for a struct or enum of modules, parameters and
9//!   constants. `#[module(skip)]` excludes a field.
10//! - `#[derive(Config)]` makes a struct a serializable configuration with a generated `new`
11//!   constructor and `with_*` setters for optional fields and fields marked `#[config(default = ...)]`.
12//! - `#[derive(RecordState)]` decomposes an optimizer or scheduler state into named tensors and
13//!   scalars for the burnpack format.
14//!
15//! The function-like macros below implement `burn::tensor::assert_shape!`,
16//! `debug_assert_shape!` and `einsum!`, and are not part of the public API.
17
18#[macro_use]
19extern crate derive_new;
20
21use proc_macro::TokenStream;
22
23pub(crate) mod config;
24pub(crate) mod einsum;
25pub(crate) mod module;
26pub(crate) mod record_state;
27pub(crate) mod shape;
28pub(crate) mod shared;
29
30/// Derive macro for the `Module` trait.
31///
32/// # Sub-modules
33///
34/// Fields are treated as sub-modules unless their type is recognized as a constant or they
35/// are marked `#[module(skip)]`. Participating fields must implement `Module` and `ModuleDisplay`.
36///
37/// Constants are detected syntactically, including `usize`, `f32`, `bool`, `String` and
38/// combinations such as `Option<usize>` and `Vec<f32>`. Unlike 0.21, unknown concrete types
39/// are not automatically skipped. Use `#[module(skip)]` for custom config/state fields that
40/// report missing `Module` / `ModuleDisplay` implementations. Do not skip trainable children.
41///
42/// ## Generics
43///
44/// Generic type parameters (e.g., `field: M`) are assumed to be sub-modules by default.
45/// If a generic field represents some other runtime state or configuration, you can use
46/// the `#[module(skip)]` attribute to provide a hint.
47///
48/// # Field Attributes
49///
50/// ## `#[module(skip)]`
51///
52/// Explicitly marks a field to be ignored by the module derive.
53///
54/// Skipped fields are excluded from parameter traversal, checkpointing, device movement and
55/// training/validation transitions. They are still cloned and included in debug display.
56/// This is equivalent to the deprecated `Ignored<T>` wrapper.
57///
58/// ### Requirements
59///
60/// The field must implement: `Debug + Clone + Send`.
61/// A `Param` field or a field using a generic explicitly bounded by `Module` cannot be skipped.
62///
63/// # Example
64///
65/// ```ignore
66/// #[derive(Module, Debug)]
67/// pub struct MyModule<M, N: NonModuleTrait> {
68///     /// A normal parameter.
69///     weights: Param<Tensor<2>>,
70///     /// A field configured at runtime.
71///     dropout_prob: f64,
72///     /// A field that is recomputed at runtime.
73///     #[module(skip)]
74///     cached_mask: Option<Tensor<2>>,
75///     /// A field that contains some debug state.
76///     debug_state: String,
77///     /// Treated as a module (default for generics).
78///     inner: M,
79///     /// Hint required: this generic is NOT a module.
80///     #[module(skip)]
81///     other: N,
82/// }
83/// ```
84#[proc_macro_derive(Module, attributes(module))]
85pub fn module_derive(input: TokenStream) -> TokenStream {
86    let input = syn::parse(input).unwrap();
87    module::derive_impl(&input)
88}
89
90/// Derive macro for the config.
91#[proc_macro_derive(Config, attributes(config))]
92pub fn config_derive(input: TokenStream) -> TokenStream {
93    let item = syn::parse(input).unwrap();
94    config::derive_impl(&item)
95}
96
97/// Derive macro for a recordable state (optimizer or learning-rate scheduler), decomposing it
98/// into named tensors and scalars for the burnpack format.
99///
100/// Supported field shapes: `Tensor<D>`, `Option<Tensor<D>>`, `Vec<Tensor<D>>`, scalars
101/// (`usize`/`isize`/`u8`..`u64`/`i8`..`i64`/`f32`/`f64`/`bool`), `Option<scalar>`, a nested
102/// `RecordState`, and `Option<Nested>`. Scalar fields must use a concrete primitive type, not a
103/// type alias (e.g. `f64`, not `LearningRate`): classification is syntactic, so an alias is treated
104/// as a nested state.
105#[proc_macro_derive(RecordState)]
106pub fn record_state_derive(input: TokenStream) -> TokenStream {
107    let input = syn::parse(input).unwrap();
108    record_state::derive_impl(&input)
109}
110
111/// Implementation of `burn_tensor::assert_shape!`. Not part of the public API.
112#[doc(hidden)]
113#[proc_macro]
114pub fn __assert_shape(input: TokenStream) -> TokenStream {
115    shape_macro(input, shape::Mode::Assert)
116}
117
118/// Implementation of `burn_tensor::debug_assert_shape!`. Not part of the public API.
119#[doc(hidden)]
120#[proc_macro]
121pub fn __debug_assert_shape(input: TokenStream) -> TokenStream {
122    shape_macro(input, shape::Mode::DebugAssert)
123}
124
125/// The input is `$crate, tensor, [pattern]`, see the wrappers in burn-tensor.
126fn shape_macro(input: TokenStream, mode: shape::Mode) -> TokenStream {
127    let input = syn::parse_macro_input!(input as shape::ShapeInput);
128    shape::expand(input, mode).into()
129}
130
131/// Implementation of `burn_tensor::einsum!`. Not part of the public API.
132#[doc(hidden)]
133#[proc_macro]
134pub fn __einsum(input: TokenStream) -> TokenStream {
135    let input = syn::parse_macro_input!(input as einsum::EinsumInput);
136    einsum::expand(input)
137        .unwrap_or_else(syn::Error::into_compile_error)
138        .into()
139}
burn_derive/lib.rs

burn_derive/
lib.rs
