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/// Bounds are inferred on each field's type: `Wrapper<T>` must implement `Module` and
45/// `ModuleDisplay`, but `T` need not implement them. Existing bounds, including custom
46/// forward traits, are preserved. Use `#[module(skip)]` for generic config/state fields;
47/// the same type parameter can also appear in a participating field.
48///
49/// Named and tuple structs support recursive children such as `Vec<Self>` and
50/// `Option<Box<Self>>`. Enums support variants containing a single unnamed field.
51///
52/// # Field Attributes
53///
54/// ## `#[module(skip)]`
55///
56/// Explicitly marks a field to be ignored by the module derive.
57///
58/// Skipped fields are excluded from parameter traversal, checkpointing, device movement and
59/// training/validation transitions. They are still cloned and included in debug display.
60/// This is equivalent to the deprecated `Ignored<T>` wrapper.
61///
62/// ### Requirements
63///
64/// The field must implement: `Debug + Clone + Send`.
65/// A `Param` field cannot be skipped.
66///
67/// # Example
68///
69/// ```ignore
70/// #[derive(Module, Debug)]
71/// pub struct MyModule<M, N: NonModuleTrait> {
72///     /// A normal parameter.
73///     weights: Param<Tensor<2>>,
74///     /// A field configured at runtime.
75///     dropout_prob: f64,
76///     /// A field that is recomputed at runtime.
77///     #[module(skip)]
78///     cached_mask: Option<Tensor<2>>,
79///     /// A field that contains some debug state.
80///     debug_state: String,
81///     /// Treated as a module (default for generics).
82///     inner: M,
83///     /// Hint required: this generic is NOT a module.
84///     #[module(skip)]
85///     other: N,
86/// }
87/// ```
88#[proc_macro_derive(Module, attributes(module))]
89pub fn module_derive(input: TokenStream) -> TokenStream {
90    let input = syn::parse(input).unwrap();
91    module::derive_impl(&input)
92}
93
94/// Derive macro for the config.
95#[proc_macro_derive(Config, attributes(config))]
96pub fn config_derive(input: TokenStream) -> TokenStream {
97    let item = syn::parse(input).unwrap();
98    config::derive_impl(&item)
99}
100
101/// Derive macro for a recordable state (optimizer or learning-rate scheduler), decomposing it
102/// into named tensors and scalars for the burnpack format.
103///
104/// Supported field shapes: `Tensor<D>`, `Option<Tensor<D>>`, `Vec<Tensor<D>>`, scalars
105/// (`usize`/`isize`/`u8`..`u64`/`i8`..`i64`/`f32`/`f64`/`bool`), `Option<scalar>`, a nested
106/// `RecordState`, and `Option<Nested>`. Scalar fields must use a concrete primitive type, not a
107/// type alias (e.g. `f64`, not `LearningRate`): classification is syntactic, so an alias is treated
108/// as a nested state.
109#[proc_macro_derive(RecordState)]
110pub fn record_state_derive(input: TokenStream) -> TokenStream {
111    let input = syn::parse(input).unwrap();
112    record_state::derive_impl(&input)
113}
114
115/// Implementation of `burn_tensor::assert_shape!`. Not part of the public API.
116#[doc(hidden)]
117#[proc_macro]
118pub fn __assert_shape(input: TokenStream) -> TokenStream {
119    shape_macro(input, shape::Mode::Assert)
120}
121
122/// Implementation of `burn_tensor::debug_assert_shape!`. Not part of the public API.
123#[doc(hidden)]
124#[proc_macro]
125pub fn __debug_assert_shape(input: TokenStream) -> TokenStream {
126    shape_macro(input, shape::Mode::DebugAssert)
127}
128
129/// The input is `$crate, tensor, [pattern]`, see the wrappers in burn-tensor.
130fn shape_macro(input: TokenStream, mode: shape::Mode) -> TokenStream {
131    let input = syn::parse_macro_input!(input as shape::ShapeInput);
132    shape::expand(input, mode).into()
133}
134
135/// Implementation of `burn_tensor::einsum!`. Not part of the public API.
136#[doc(hidden)]
137#[proc_macro]
138pub fn __einsum(input: TokenStream) -> TokenStream {
139    let input = syn::parse_macro_input!(input as einsum::EinsumInput);
140    einsum::expand(input)
141        .unwrap_or_else(syn::Error::into_compile_error)
142        .into()
143}
burn_derive/lib.rs

burn_derive/
lib.rs
