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