burn_derive/
lib.rs

1#![warn(missing_docs)]
2
3//! The derive crate of Burn.
4
5#[macro_use]
6extern crate derive_new;
7
8use proc_macro::TokenStream;
9
10pub(crate) mod config;
11pub(crate) mod einsum;
12pub(crate) mod module;
13pub(crate) mod record_state;
14pub(crate) mod shape;
15pub(crate) mod shared;
16
17/// Derive macro for the `Module` trait.
18///
19/// # Sub-modules
20///
21/// By default, the macro automatically detects sub-modules and parameters as module types.
22///
23/// Any field not recognized as a module type is assumed to be a non-module
24/// and is skipped by the module system (not persistent, not visited).
25///
26/// ## Generics
27///
28/// Generic type parameters (e.g., `field: M`) are assumed to be sub-modules by default.
29/// If a generic field represents some other runtime state or configuration, you can use
30/// the `#[module(skip)]` attribute to provide a hint.
31///
32/// # Field Attributes
33///
34/// ## `#[module(skip)]`
35///
36/// Explicitly marks a field to be ignored by the module derive.
37///
38/// Skipped fields are not parameters, not modules, and are not persistent.
39/// This is equivalent to the deprecated `Ignored<T>` wrapper.
40///
41/// ### Requirements
42///
43/// The field must implement: `Debug + Clone + Send`.
44///
45/// # Example
46///
47/// ```ignore
48/// #[derive(Module, Debug)]
49/// pub struct MyModule<M, N: NonModuleTrait> {
50///     /// A normal parameter.
51///     weights: Param<Tensor<2>>,
52///     /// A field configured at runtime.
53///     dropout_prob: f64,
54///     /// A field that is recomputed at runtime.
55///     cached_mask: Option<Tensor<2>>,
56///     /// A field that contains some debug state.
57///     debug_state: String,
58///     /// Treated as a module (default for generics).
59///     inner: M,
60///     /// Hint required: this generic is NOT a module.
61///     #[module(skip)]
62///     other: N,
63/// }
64/// ```
65#[proc_macro_derive(Module, attributes(module))]
66pub fn module_derive(input: TokenStream) -> TokenStream {
67    let input = syn::parse(input).unwrap();
68    module::derive_impl(&input)
69}
70
71/// Derive macro for the config.
72#[proc_macro_derive(Config, attributes(config))]
73pub fn config_derive(input: TokenStream) -> TokenStream {
74    let item = syn::parse(input).unwrap();
75    config::derive_impl(&item)
76}
77
78/// Derive macro for a recordable state (optimizer or learning-rate scheduler), decomposing it
79/// into named tensors and scalars for the burnpack format.
80///
81/// Supported field shapes: `Tensor<D>`, `Option<Tensor<D>>`, `Vec<Tensor<D>>`, scalars
82/// (`usize`/`isize`/`u8`..`u64`/`i8`..`i64`/`f32`/`f64`/`bool`), `Option<scalar>`, a nested
83/// `RecordState`, and `Option<Nested>`. Scalar fields must use a concrete primitive type, not a
84/// type alias (e.g. `f64`, not `LearningRate`): classification is syntactic, so an alias is treated
85/// as a nested state.
86#[proc_macro_derive(RecordState)]
87pub fn record_state_derive(input: TokenStream) -> TokenStream {
88    let input = syn::parse(input).unwrap();
89    record_state::derive_impl(&input)
90}
91
92/// Implementation of `burn_tensor::assert_shape!`. Not part of the public API.
93#[doc(hidden)]
94#[proc_macro]
95pub fn __assert_shape(input: TokenStream) -> TokenStream {
96    shape_macro(input, shape::Mode::Assert)
97}
98
99/// Implementation of `burn_tensor::debug_assert_shape!`. Not part of the public API.
100#[doc(hidden)]
101#[proc_macro]
102pub fn __debug_assert_shape(input: TokenStream) -> TokenStream {
103    shape_macro(input, shape::Mode::DebugAssert)
104}
105
106/// The input is `$crate, tensor, [pattern]`, see the wrappers in burn-tensor.
107fn shape_macro(input: TokenStream, mode: shape::Mode) -> TokenStream {
108    let input = syn::parse_macro_input!(input as shape::ShapeInput);
109    shape::expand(input, mode).into()
110}
111
112/// Implementation of `burn_tensor::einsum!`. Not part of the public API.
113#[doc(hidden)]
114#[proc_macro]
115pub fn __einsum(input: TokenStream) -> TokenStream {
116    let input = syn::parse_macro_input!(input as einsum::EinsumInput);
117    einsum::expand(input)
118        .unwrap_or_else(syn::Error::into_compile_error)
119        .into()
120}
burn_derive/lib.rs

burn_derive/
lib.rs
