burn_core/module/param/base.rs
1#![allow(clippy::bool_assert_comparison)]
2
3use crate::module::LoraAdapter;
4use crate::module::Reparameterization;
5
6use super::ParamId;
7use super::reparameterization_dyn::DynReparameterization;
8use super::sync_once_cell::SyncOnceCell;
9use alloc::format;
10
11use alloc::boxed::Box;
12use burn_std::sync::{Arc, RwLock};
13use burn_tensor::{Device, Shape};
14use core::ops::Deref;
15
16#[cfg(target_has_atomic = "ptr")]
17type Mapper<T> = Arc<dyn Fn(T) -> T + Send + Sync>;
18
19#[cfg(not(target_has_atomic = "ptr"))]
20type Mapper<T> = Arc<Box<dyn Fn(T) -> T + Send + Sync>>;
21
22#[cfg(target_has_atomic = "ptr")]
23fn new_mapper<T, F: Fn(T) -> T + Send + Sync + 'static>(func: F) -> Mapper<T> {
24 Arc::new(func)
25}
26
27#[cfg(not(target_has_atomic = "ptr"))]
28fn new_mapper<T, F: Fn(T) -> T + Send + Sync + 'static>(func: F) -> Mapper<T> {
29 Arc::new(Box::new(func))
30}
31
32type InitFn<P> = Box<dyn FnOnce(&Device, bool) -> P + Send + Sync>;
33
34fn new_init_fn<P: ParameterValue, F: FnOnce(&Device, bool) -> P + Send + Sync + 'static>(
35 func: F,
36) -> InitFn<P> {
37 Box::new(func)
38}
39
40/// Coordinates lazy initialization across all clones of a [`Param`].
41///
42/// The sole purpose of this shared state is to ensure the initialization function runs at most
43/// once: whichever clone first calls [`val`](Param::val) initializes the value, and all other
44/// clones observe the same result.
45///
46/// # State Management
47///
48/// **Two logical states:**
49///
50/// 1. **Initialized**: `value` contains the parameter value and `initialization` is `None`.
51/// 2. **Lazily Managed**: `initialization` contains `Some(RwLock<...>)`.
52/// - *Before initialization*: `value` is empty, inner option is `Some(Uninitialized<T>)`.
53/// - *After initialization*: `value` contains the parameter value, inner option is `None`.
54///
55/// The transition from uninitialized to initialized happens exactly once and is synchronized
56/// across all clones.
57pub(crate) struct LazyInitState<T: ParameterValue> {
58 /// The SyncOnceCell holding the initialized parameter value.
59 /// Empty for uninitialized parameters, populated after first access or explicit initialization.
60 pub value: SyncOnceCell<T>,
61 /// The deferred initialization state for lazy parameters.
62 ///
63 /// **State Transitions:**
64 /// - Initialized params: `None`
65 /// - Uninitialized params: `Some(RwLock<Some(Uninitialized<T>)>)`
66 /// - After lazy init triggers: `Some(RwLock<None>)` (inner Option is taken)
67 pub initialization: Option<RwLock<Option<Uninitialized<T>>>>,
68}
69
70impl<T: ParameterValue> LazyInitState<T> {
71 /// Create a new parameter state that is already initialized.
72 fn initialized(value: T) -> Arc<Self> {
73 Arc::new(Self {
74 value: SyncOnceCell::initialized(value),
75 initialization: None,
76 })
77 }
78
79 /// Create a new parameter state that is not already initialized.
80 fn uninitialized(uninit: Uninitialized<T>) -> Arc<Self> {
81 Arc::new(Self {
82 value: SyncOnceCell::new(),
83 initialization: Some(RwLock::new(Some(uninit))),
84 })
85 }
86
87 /// Gets the parameter value, initializing it lazily if needed.
88 fn val(&self) -> &T {
89 self.value.get_or_init(|| {
90 let mut init = self
91 .initialization
92 .as_ref()
93 .expect("Should have an initialization when no state provided.")
94 .write();
95 let state = init.take().expect("Should exist when not initialized");
96 state.initialize()
97 })
98 }
99}
100
101/// Parameters are identified, traversable values owned by [modules](crate::module::Module).
102///
103/// Most parameters contain [tensors](crate::tensor::Tensor) that can be updated during training
104/// and loaded during inference. Parameters can also contain module control state such as
105/// [`Flag`](crate::module::Flag), allowing it to participate in module traversal and parameter
106/// group selection without pretending to be a tensor.
107///
108/// # Cloning
109///
110/// Cloning a parameter is always cheap; it never allocates or initializes tensors.
111/// Clones share the same lazy initialization state, so initialization happens at most once and
112/// all clones resolve to the same value regardless of which one triggered it.
113///
114/// This sharing is strictly scoped to lazy initialization. It only guarantees that all clones
115/// observe the same initialization result. Subsequent transformations operate on independent
116/// parameter values and never propagate across clones.
117pub struct Param<T: ParameterValue> {
118 /// The unique ID of this parameter. This is used by eg. optimizers to associate a gradient with a specific parameter.
119 pub id: ParamId,
120 /// Shared lazy initialization state across all clones of this parameter.
121 /// The `Arc` exists solely to coordinate lazy initialization. It is not a general
122 /// shared-ownership mechanism. Any mutation forks into a new `LazyInitState`.
123 pub(crate) state: Arc<LazyInitState<T>>,
124 pub(crate) param_mapper: ParamMapper<T>,
125 /// Whether this value is configured to participate in training behavior.
126 ///
127 /// This is authoritative and kept separately from the effective value so transformations and
128 /// backends that cannot currently activate the value don't erase the state that
129 /// [`Module::valid`](crate::module::Module::valid) temporarily suspends and
130 /// `train` restores.
131 pub(crate) is_active: bool,
132 /// Optional transformation that materializes the effective value from the stored base.
133 pub(crate) reparameterization: Option<Box<dyn DynReparameterization>>,
134}
135
136#[derive(Clone)]
137/// Applies transformations when loading and saving parameters.
138///
139/// # Mapper System
140///
141/// `ParamMapper<T>` allows applying transformations during serialization and deserialization:
142/// - `load: Option<Mapper<T>>` - transformation during deserialization (applied in `transform_for_load()`)
143/// - `save: Option<Mapper<T>>` - transformation during serialization (applied in `transform_for_save()`)
144///
145/// These are commonly used for:
146/// - Quantization/dequantization
147/// - Precision conversion (e.g., FP32 ↔ FP16)
148/// - Custom parameter transformations
149pub struct ParamMapper<T: ParameterValue> {
150 load: Option<Mapper<T>>,
151 save: Option<Mapper<T>>,
152 /// Carries configured training state across `consume` / `from_mapped_value` reconstruction.
153 mapped_is_active: Option<bool>,
154}
155
156impl<T: ParameterValue> core::fmt::Debug for ParamMapper<T> {
157 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
158 f.write_fmt(format_args!(
159 "ParamMapper {{ load: {}, save: {} }}",
160 self.load.is_some(),
161 self.save.is_some(),
162 ))
163 }
164}
165
166impl<T: ParameterValue> ParamMapper<T> {
167 /// Applies the transformation when loading the given parameter.
168 pub fn on_load(&self, param: T) -> T {
169 match &self.load {
170 Some(mapper) => mapper(param),
171 None => param,
172 }
173 }
174 /// Applies the transformation when saving the given parameter.
175 pub fn on_save(&self, param: T) -> T {
176 match &self.save {
177 Some(mapper) => mapper(param),
178 None => param,
179 }
180 }
181}
182
183impl<T: ParameterValue> Default for ParamMapper<T> {
184 fn default() -> Self {
185 Self {
186 load: None,
187 save: None,
188 mapped_is_active: None,
189 }
190 }
191}
192
193impl<T: ParameterValue> core::fmt::Display for Param<T> {
194 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
195 f.write_str(format!("Param: {}", self.id).as_str())
196 }
197}
198
199impl<T: ParameterValue> core::fmt::Debug for Param<T> {
200 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
201 f.write_str(format!("Param: {} - {:?}", self.id, self.param_mapper).as_str())
202 }
203}
204
205pub(crate) mod sealed {
206 use super::DynReparameterization;
207
208 pub trait Sealed: Sized {
209 /// Whether this value currently participates in the module's training behavior.
210 fn is_active(&self) -> bool;
211
212 /// Materialize a parameter with an attached reparameterization.
213 ///
214 /// # Notes
215 /// This is part of the sealed trait to avoid [`DynReparameterization`] from showing up in the
216 /// public `ParameterValue` trait.
217 fn materialize(self, reparameterization: &dyn DynReparameterization) -> Self {
218 let _ = reparameterization;
219 self
220 }
221 }
222}
223
224/// A value that can be identified and traversed inside a [`Param`].
225///
226/// # Notes
227/// This trait is intentionally sealed to keep the set of parameters closed.
228///
229/// Although exposed publicly, parameter values are not meant to be extensible. The closed set
230/// includes device-backed values such as [`Tensor`](crate::Tensor) and device-independent control
231/// values such as [`Flag`](crate::module::Flag).
232pub trait ParameterValue: sealed::Sealed + Clone + core::fmt::Debug + Send {}
233
234/// A tensor-backed [`ParameterValue`] with autodiff, device and loading capabilities.
235///
236/// General module-owned values can implement [`ParameterValue`] without pretending to require
237/// gradients, reside on a device or have a tensor shape.
238pub trait Parameter: ParameterValue {
239 /// Fetch the gradient requirement.
240 fn is_require_grad(&self) -> bool;
241
242 /// Set the gradient requirement.
243 ///
244 /// Float tensors without autodiff remain unchanged. [`Param::set_require_grad`] stores
245 /// the configured training state so it can be applied when entering training mode.
246 fn set_require_grad(self, require_grad: bool) -> Self;
247
248 /// Fetch the device.
249 fn device(&self) -> Device;
250
251 /// Fetch the shape of the parameter.
252 fn shape(&self) -> Shape;
253
254 /// Moves the parameter to the target device if it is not already on it,
255 /// applying any kind-specific preparation required for the loading lifecycle (e.g. detach).
256 fn load_to_device(self, device: &Device) -> Self;
257}
258
259/// The deferred initialization state for lazy parameters.
260#[allow(clippy::type_complexity)]
261pub(crate) struct Uninitialized<P: ParameterValue> {
262 /// The initialization function. Called with `(device, is_require_grad) -> Parameter`.
263 init: InitFn<P>,
264 /// The target device on which the parameter should be initialized.
265 /// Used by `lazy_device()` to provide device information without triggering initialization.
266 pub(crate) device: Device,
267 /// The gradient requirement for the parameter.
268 /// Used by `lazy_is_require_grad()` to provide gradient settings without triggering initialization.
269 pub(crate) is_require_grad: bool,
270 /// The shape of the tensor parameter.
271 /// Used by `lazy_shape()` to provide shape information without triggering initialization.
272 pub(crate) shape: Shape,
273}
274
275impl<P: ParameterValue> Uninitialized<P> {
276 /// Runs the initialization function.
277 ///
278 /// This is called by [Param::val] when accessing an uninitialized parameter for the first time.
279 /// The function receives the stored device and effective gradient requirement. On a plain
280 /// device the requirement is false; the configured training state is preserved for `train()`.
281 fn initialize(self) -> P {
282 (self.init)(
283 &self.device,
284 self.is_require_grad && self.device.is_autodiff(),
285 )
286 }
287}
288
289impl<T: ParameterValue> Param<T> {
290 /// Create a new parameter that is already initialized.
291 pub fn initialized(id: ParamId, value: T) -> Self {
292 let is_active = value.is_active();
293 Self {
294 id,
295 state: LazyInitState::initialized(value),
296 param_mapper: Default::default(),
297 is_active,
298 reparameterization: None,
299 }
300 }
301
302 /// Gets the effective parameter value, initializing it lazily if needed.
303 ///
304 /// For initialized parameters, this returns a clone of the cached value.
305 /// For uninitialized parameters, this triggers initialization.
306 ///
307 /// When a reparameterization is attached, this materializes its effective value. Use
308 /// [`base`](Self::base) to access the raw stored value without materialization. Conceptually,
309 /// materialization composes the structural base with the attached reparameterization state.
310 pub fn val(&self) -> T {
311 let base = self.deref().clone();
312 match &self.reparameterization {
313 Some(reparameterization) => base.materialize(reparameterization.as_ref()),
314 None => base,
315 }
316 }
317
318 /// Gets the raw stored parameter value **without** applying its reparameterization.
319 pub fn base(&self) -> T {
320 self.deref().clone()
321 }
322
323 /// Check if the parameter has been initialized.
324 ///
325 /// Returns `true` if the parameter's value has been computed and cached,
326 /// `false` if it's still lazy and will be initialized on first access.
327 pub fn is_initialized(&self) -> bool {
328 self.state.value.get().is_some()
329 }
330
331 /// Gets the parameter's value while consuming the parameter.
332 pub fn into_value(self) -> T {
333 self.consume().1
334 }
335
336 /// Gets the parameter id and raw value while consuming the parameter.
337 ///
338 /// Any attached reparameterization is dropped. During module traversal,
339 /// `Param<Tensor<D>>::map` detaches the reparameterization before calling
340 /// `ModuleMapper::map_float`, traverses it separately, and reattaches it afterward.
341 pub fn consume(self) -> (ParamId, T, ParamMapper<T>) {
342 let tensor = self.deref().clone();
343 let mut param_mapper = self.param_mapper;
344 param_mapper.mapped_is_active = Some(self.is_active);
345
346 core::mem::drop(self.state);
347
348 (self.id, tensor, param_mapper)
349 }
350
351 /// Execute the given function on the inner value while preserving its configured training
352 /// state.
353 ///
354 /// Transforming the effective value must not implicitly change what
355 /// [`Module::train`](crate::module::Module::train) restores. Use the explicit activation APIs,
356 /// such as [`Param::set_require_grad`] for tensor parameters, to change that state.
357 pub fn map<F: FnOnce(T) -> T>(self, func: F) -> Self {
358 let (id, tensor, param_mapper) = self.consume();
359 let tensor = func(tensor);
360 Self::from_mapped_value(id, tensor, param_mapper)
361 }
362
363 /// Create an initialized parameter with the given id, value, and param mapper.
364 ///
365 /// This is a helper method for creating parameters while preserving the param mapper,
366 /// typically used in ModuleMapper implementations.
367 pub fn from_mapped_value(id: ParamId, value: T, mut param_mapper: ParamMapper<T>) -> Self {
368 let is_active = param_mapper
369 .mapped_is_active
370 .take()
371 .unwrap_or_else(|| value.is_active());
372 Self {
373 id,
374 state: LazyInitState::initialized(value),
375 param_mapper,
376 is_active,
377 reparameterization: None,
378 }
379 }
380}
381
382impl<T: Parameter> Param<T> {
383 /// Create a new parameter that is not already initialized.
384 ///
385 /// The initializer receives the device and effective gradient requirement, which is false
386 /// on devices without autodiff. The requested `is_require_grad` setting is preserved for
387 /// [`Module::train`](crate::module::Module::train). The initializer must create the value on
388 /// the device it receives, which is not `device` when the parameter moves before initializing.
389 pub fn uninitialized<F>(
390 id: ParamId,
391 init: F,
392 device: Device,
393 is_require_grad: bool,
394 shape: Shape,
395 ) -> Self
396 where
397 F: FnOnce(&Device, bool) -> T + Send + Sync + 'static,
398 {
399 Self {
400 id,
401 state: LazyInitState::uninitialized(Uninitialized {
402 init: new_init_fn(init),
403 device,
404 is_require_grad,
405 shape,
406 }),
407 param_mapper: Default::default(),
408 is_active: is_require_grad,
409 reparameterization: None,
410 }
411 }
412
413 pub(crate) fn reparameterization_dyn(&self) -> Option<&dyn DynReparameterization> {
414 self.reparameterization.as_deref()
415 }
416
417 /// The concrete [reparameterization](Reparameterization) attached to this parameter, if any.
418 pub fn reparameterization<R: Reparameterization>(&self) -> Option<&R> {
419 self.reparameterization_dyn()?.as_any().downcast_ref()
420 }
421
422 /// The LoRA [adapter](LoraAdapter) attached to this parameter, if any.
423 pub fn adapter(&self) -> Option<&LoraAdapter> {
424 self.reparameterization()
425 }
426
427 /// Returns a cheap clone of this parameter with its reparameterization detached.
428 ///
429 /// The clone shares the same lazy-initialization state, so the raw base value is not
430 /// duplicated. Used to route the optimizer/record traversal over the structural base.
431 pub(crate) fn without_reparameterization(&self) -> Self {
432 Self {
433 id: self.id,
434 state: self.state.clone(),
435 param_mapper: self.param_mapper.clone(),
436 is_active: self.is_active,
437 reparameterization: None,
438 }
439 }
440
441 pub(crate) fn with_dyn_reparameterization(
442 mut self,
443 reparameterization: Option<Box<dyn DynReparameterization>>,
444 ) -> Self {
445 self.reparameterization = reparameterization;
446 self
447 }
448
449 /// Runs a transformation on the parameter when loading.
450 pub fn load_mapper<F: Fn(T) -> T + Send + Sync + 'static>(mut self, func: F) -> Self {
451 self.param_mapper.load = Some(new_mapper(func));
452
453 self
454 }
455
456 /// Runs a transformation on the parameter when saving.
457 pub fn save_mapper<F: Fn(T) -> T + Send + Sync + 'static>(mut self, func: F) -> Self {
458 self.param_mapper.save = Some(new_mapper(func));
459
460 self
461 }
462
463 /// Returns a new parameter whose initialization value is transformed by the given function.
464 ///
465 /// If the parameter is still uninitialized (lazy), the transformation is chained onto the
466 /// existing initialization without triggering evaluation. If the parameter is already
467 /// initialized, it immediately applies the transformation to the current value.
468 pub fn init_mapper<F: Fn(T) -> T + Send + Sync + 'static>(self, func: F) -> Self
469 where
470 T: Sync + 'static,
471 {
472 let initialization = match &self.state.initialization {
473 Some(init) => init,
474 None => return self.map(func),
475 };
476
477 let mut init = initialization.write();
478
479 match init.as_mut() {
480 Some(value) => {
481 let device = value.device.clone();
482 let is_require_grad = value.is_require_grad;
483 let shape = value.shape.clone();
484 core::mem::drop(init);
485
486 let mut base = self;
487 Self {
488 id: base.id,
489 param_mapper: base.param_mapper.clone(),
490 is_active: base.is_active,
491 reparameterization: None,
492 state: LazyInitState::uninitialized(Uninitialized {
493 // A clone sharing `base` still initializes it where it was built. `func`
494 // maps an untracked value: mapped from a tracked leaf, it would be a
495 // non-leaf that can't require grad.
496 init: new_init_fn(move |device, require_grad| {
497 // Move the base and attached state together before materializing.
498 // Detaching here prevents the map_to_device fallback from dropping it.
499 let reparameterization = base.reparameterization.take();
500 let base = base
501 .map_to_device(device, |value| value.load_to_device(device))
502 .with_dyn_reparameterization(
503 reparameterization.map(|state| state.to_device_dyn(device)),
504 );
505 let value = base.val().set_require_grad(false);
506 func(value).set_require_grad(require_grad)
507 }),
508 device,
509 is_require_grad,
510 shape,
511 }),
512 }
513 }
514 None => {
515 core::mem::drop(init);
516 self.map(func)
517 }
518 }
519 }
520
521 /// The device on which the parameter is or will be initialized, **without triggering initialization**.
522 ///
523 /// This is critical for the load optimization: when loading tensors into an uninitialized parameter,
524 /// we need to know the target device to move the loaded tensor appropriately, but we don't want to
525 /// trigger the initialization function (which would allocate an unnecessary tensor).
526 ///
527 /// Use this instead of [crate::tensor::Tensor::device] when you need the device but want to
528 /// preserve lazy initialization.
529 pub fn lazy_device(&self) -> Device {
530 let initialization = match &self.state.initialization {
531 Some(init) => init,
532 None => return self.device(),
533 };
534
535 let init = initialization.read();
536
537 match init.as_ref() {
538 Some(value) => value.device.clone(),
539 None => self.device(),
540 }
541 }
542
543 /// Put the parameter on `device`: one not initialized yet initializes there, unless a clone
544 /// shares it, and any other is initialized if needed and moved with `move_value`. Either way it
545 /// keeps the autodiff context it was built with, as moving a value does.
546 pub(crate) fn map_to_device(
547 mut self,
548 device: &Device,
549 move_value: impl FnOnce(T) -> T,
550 ) -> Self {
551 let retargeted = Arc::get_mut(&mut self.state)
552 .and_then(|state| state.initialization.as_ref())
553 .is_some_and(|initialization| match initialization.write().as_mut() {
554 Some(uninitialized) => {
555 uninitialized.device = device
556 .clone()
557 .with_autodiff_context_from(&uninitialized.device);
558 true
559 }
560 None => false,
561 });
562
563 if retargeted {
564 self
565 } else {
566 self.map(move_value)
567 }
568 }
569
570 /// The gradient requirement on which the parameter is or will be initialized, **without triggering initialization**.
571 ///
572 /// Similar to [lazy_device](Self::lazy_device), this is critical for the load optimization.
573 /// When loading tensors into an uninitialized parameter, we need to apply the correct gradient
574 /// setting to the loaded tensor without triggering the initialization function.
575 ///
576 /// # Notes
577 ///
578 /// This is a crate-private function, since users are not expected to use `is_require_grad` of an
579 /// uninitialized module to then override its value. All low-level functions should be provided
580 /// by `burn` and should handle those details.
581 pub(crate) fn lazy_is_require_grad(&self) -> bool {
582 let initialization = match &self.state.initialization {
583 Some(init) => init,
584 None => return self.is_active,
585 };
586
587 let init = initialization.read();
588
589 match init.as_ref() {
590 Some(value) => value.is_require_grad,
591 None => self.is_active,
592 }
593 }
594
595 /// Override the gradient-tracking setting for the current parameter.
596 ///
597 /// On a backend without autodiff the effective tensor remains detached, but the setting is
598 /// preserved and takes effect through [`Module::train`](crate::module::Module::train).
599 pub fn set_require_grad(mut self, require_grad: bool) -> Self {
600 let initialization = match &self.state.initialization {
601 Some(init) => init,
602 None => {
603 let mut param = self.map(|tensor| tensor.set_require_grad(require_grad));
604 param.is_active = require_grad;
605 return param;
606 }
607 };
608
609 let mut init = initialization.write();
610 let mut is_lazy = false;
611
612 if let Some(value) = init.as_mut() {
613 is_lazy = true;
614 value.is_require_grad = require_grad;
615 };
616
617 core::mem::drop(init);
618
619 if is_lazy {
620 self.is_active = require_grad;
621 return self;
622 }
623
624 let mut param = self.map(|tensor| tensor.set_require_grad(require_grad));
625 param.is_active = require_grad;
626 param
627 }
628
629 /// The shape of the parameter, **without triggering initialization**.
630 ///
631 /// This is critical for shape validation during loading: when applying tensors to an
632 /// uninitialized parameter, we need to validate the shape without triggering the
633 /// initialization function (which would allocate an unnecessary tensor).
634 ///
635 /// Use this instead of [crate::tensor::Tensor::shape] when you need the shape but want to
636 /// preserve lazy initialization.
637 pub fn lazy_shape(&self) -> burn_tensor::Shape {
638 let initialization = match &self.state.initialization {
639 Some(init) => init,
640 None => return self.shape(),
641 };
642
643 let init = initialization.read();
644
645 match init.as_ref() {
646 Some(value) => value.shape.clone(),
647 None => self.shape(),
648 }
649 }
650
651 /// Transform a parameter for loading by applying load transformations.
652 ///
653 /// This method is used to restore a parameter from a tensor (typically during deserialization).
654 /// It ensures the tensor is moved to the expected device, applies the param mapper's
655 /// `on_load` transformation, and preserves the autodiff settings (require_grad).
656 pub fn transform_for_load(self, tensor: T, param_id: ParamId) -> Self {
657 let mut new_tensor = tensor;
658
659 let mapper = self.param_mapper.clone();
660
661 let expected_device = self.lazy_device();
662 let expected_require_grad = self.lazy_is_require_grad();
663
664 // Make sure we load the tensor into the same module device.
665 new_tensor = new_tensor.load_to_device(&expected_device);
666
667 new_tensor = mapper.on_load(new_tensor);
668
669 // Make sure we load the tensor with the same autodiff setting.
670 new_tensor = new_tensor.set_require_grad(expected_require_grad);
671
672 let mut loaded = Self::initialized(param_id, new_tensor);
673 loaded.param_mapper = mapper;
674 loaded.is_active = expected_require_grad;
675 loaded
676 }
677
678 /// Transform a parameter for saving by applying save transformations.
679 ///
680 /// This method is used to prepare a parameter for saving (typically during serialization).
681 /// It applies the param mapper's `on_save` transformation, which can be used
682 /// to modify the tensor before serialization (e.g., quantization, precision conversion).
683 pub fn transform_for_save(&self) -> Self {
684 let mut tensor = self.val();
685 let mapper = self.param_mapper.clone();
686
687 tensor = mapper.on_save(tensor);
688
689 let mut transformed = Self::initialized(self.id, tensor);
690 transformed.is_active = self.is_active;
691 transformed
692 }
693}
694
695impl<T: ParameterValue> Clone for Param<T> {
696 fn clone(&self) -> Self {
697 Self {
698 id: self.id,
699 state: self.state.clone(),
700 param_mapper: self.param_mapper.clone(),
701 is_active: self.is_active,
702 reparameterization: self.reparameterization.clone(),
703 }
704 }
705}
706
707impl<T: ParameterValue> Deref for Param<T> {
708 type Target = T;
709
710 fn deref(&self) -> &Self::Target {
711 self.state.val()
712 }
713}
714#[cfg(test)]
715mod tests {
716 use super::*;
717 use burn_tensor::Tensor;
718
719 // Param<T> should be Sync so that models can be shared across threads
720 // (e.g. parallel inference with rayon).
721 fn _assert_sync<T: Sync>() {}
722 fn _assert_parameter_value<T: ParameterValue>() {}
723 fn _assert_parameter<T: Parameter>() {}
724
725 #[test]
726 fn param_is_sync() {
727 fn check() {
728 _assert_sync::<Param<Tensor<2>>>();
729 }
730 check();
731 }
732
733 #[test]
734 fn parameter_capabilities_are_split_by_value_kind() {
735 _assert_parameter_value::<crate::module::Flag>();
736 _assert_parameter_value::<Tensor<1>>();
737 _assert_parameter::<Tensor<1>>();
738 }
739
740 /// Concurrent lazy initialization must not panic.
741 ///
742 /// Multiple threads call `val()` on an uninitialized `Param` simultaneously.
743 /// `SyncOnceCell::get_or_init` guarantees only one thread runs the initializer;
744 /// the others block and receive the same value.
745 #[cfg(feature = "std")]
746 #[test]
747 fn param_concurrent_lazy_init() {
748 use alloc::vec::Vec;
749
750 let device = Default::default();
751
752 let param: Param<Tensor<2>> = Param::uninitialized(
753 ParamId::new(),
754 |device, _require_grad| Tensor::random([2, 3], Default::default(), device),
755 device,
756 false,
757 [2, 3].into(),
758 );
759
760 // Share across threads via ¶m (requires Sync).
761 std::thread::scope(|s| {
762 let handles: Vec<_> = (0..4).map(|_| s.spawn(|| param.val())).collect();
763
764 let results: Vec<_> = handles.into_iter().map(|h| h.join().unwrap()).collect();
765
766 // All threads must get the same value.
767 let expected = results[0].to_data();
768 for result in &results[1..] {
769 assert_eq!(result.to_data(), expected);
770 }
771 });
772 }
773
774 #[test]
775 fn param_clones_share_lazy_initialization() {
776 let device = Default::default();
777
778 // We use random values so that if it initializes twice, the data will mismatch.
779 let param_original: Param<Tensor<2>> = Param::uninitialized(
780 ParamId::new(),
781 |device, _require_grad| Tensor::random([2, 3], Default::default(), device),
782 device,
783 false,
784 [2, 3].into(),
785 );
786
787 // Regression: https://github.com/tracel-ai/burn/issues/5040
788 // Clone the parameter while it is still uninitialized.
789 // Previously, this would clone the init function only, leading to different parameter states.
790 let param_clone = param_original.clone();
791
792 let tensor_original = param_original.val();
793 assert!(param_original.is_initialized());
794 assert!(param_clone.is_initialized());
795
796 let tensor_clone = param_clone.val();
797
798 tensor_original
799 .into_data()
800 .assert_eq(&tensor_clone.into_data(), true);
801 }
802
803 #[test]
804 fn param_set_require_grad_forks_from_shared_state() {
805 let device = Default::default();
806
807 let param1: Param<Tensor<2>> = Param::uninitialized(
808 ParamId::new(),
809 |device, require_grad| Tensor::ones([2, 3], device).set_require_grad(require_grad),
810 device,
811 true,
812 [2, 3].into(),
813 );
814
815 // Clone param; both now point to the exact same Arc<LazyInitState>
816 let param2 = param1.clone();
817
818 // Force initialization via the first clone.
819 let _tensor1 = param1.val();
820 assert!(param1.is_initialized());
821 assert!(param2.is_initialized());
822
823 // set_require_grad intentionally forks: param2 gets a new Arc with the mutated tensor.
824 let param2 = param2.set_require_grad(false);
825
826 // The fork produced the correct require_grad state.
827 assert_eq!(param2.is_active, false);
828 assert_eq!(param1.is_active, true); // param1 is unaffected
829
830 // Values are still identical (same tensor data, different grad setting).
831 param1
832 .val()
833 .into_data()
834 .assert_eq(¶m2.val().into_data(), true);
835 }
836}