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