burn_core/module/param/
tensor.rs

1use super::reparameterization_dyn::{self, DynReparameterization};
2use super::{Param, ParamId, Parameter, ParameterValue, Reparameterization};
3use crate::module::{
4    Content, Module, ModuleDisplay, ModuleDisplayDefault, ModuleMapper, ModuleVisitor,
5};
6use alloc::{boxed::Box, format, string::ToString, vec::Vec};
7use burn_tensor::{Bool, Device, Float, Int, Tensor, TensorData};
8
9impl<const D: usize> super::sealed::Sealed for Tensor<D, Float> {
10    fn is_active(&self) -> bool {
11        Tensor::is_require_grad(self)
12    }
13
14    fn apply_reparameterization(self, reparameterization: &dyn DynReparameterization) -> Self {
15        *reparameterization
16            .apply_dyn(Box::new(self))
17            .downcast::<Tensor<D>>()
18            .expect("Reparameterization should preserve tensor rank")
19    }
20}
21impl<const D: usize> super::sealed::Sealed for Tensor<D, Int> {
22    fn is_active(&self) -> bool {
23        false
24    }
25}
26impl<const D: usize> super::sealed::Sealed for Tensor<D, Bool> {
27    fn is_active(&self) -> bool {
28        false
29    }
30}
31
32impl<const D: usize> ParameterValue for Tensor<D, Float> {}
33
34impl<const D: usize> Parameter for Tensor<D, Float> {
35    fn is_require_grad(&self) -> bool {
36        Tensor::is_require_grad(self)
37    }
38
39    fn set_require_grad(self, require_grad: bool) -> Self {
40        // Parameters keep their configured training state separately from the effective tensor.
41        if require_grad && !self.is_autodiff() {
42            self
43        } else {
44            Tensor::set_require_grad(self, require_grad)
45        }
46    }
47
48    fn device(&self) -> Device {
49        Tensor::device(self)
50    }
51
52    fn shape(&self) -> burn_std::Shape {
53        Tensor::shape(self)
54    }
55
56    fn load_to_device(self, device: &Device) -> Self {
57        if self.device() != *device {
58            Tensor::to_device(self, device).detach()
59        } else {
60            self
61        }
62    }
63}
64
65impl<const D: usize> ParameterValue for Tensor<D, Int> {}
66
67impl<const D: usize> Parameter for Tensor<D, Int> {
68    fn is_require_grad(&self) -> bool {
69        false
70    }
71
72    fn set_require_grad(self, _require_grad: bool) -> Self {
73        self
74    }
75
76    fn device(&self) -> Device {
77        Tensor::device(self)
78    }
79
80    fn shape(&self) -> burn_std::Shape {
81        Tensor::shape(self)
82    }
83
84    fn load_to_device(self, device: &Device) -> Self {
85        if self.device() != *device {
86            Tensor::to_device(self, device)
87        } else {
88            self
89        }
90    }
91}
92
93impl<const D: usize> ParameterValue for Tensor<D, Bool> {}
94
95impl<const D: usize> Parameter for Tensor<D, Bool> {
96    fn is_require_grad(&self) -> bool {
97        false
98    }
99
100    fn set_require_grad(self, _require_grad: bool) -> Self {
101        self
102    }
103
104    fn device(&self) -> Device {
105        Tensor::device(self)
106    }
107
108    fn shape(&self) -> burn_std::Shape {
109        Tensor::shape(self)
110    }
111
112    fn load_to_device(self, device: &Device) -> Self {
113        if self.device() != *device {
114            Tensor::to_device(self, device)
115        } else {
116            self
117        }
118    }
119}
120
121impl<const D: usize> Param<Tensor<D>> {
122    /// Create a new trainable parameter from a float tensor.
123    ///
124    /// # Warnings
125    ///
126    /// We strongly recommend using [Param::uninitialized] if you are using this method to
127    /// initialize parameters inside a module, since the tensor initialization will be lazy,
128    /// making the loading of weights more performant.
129    pub fn from_tensor(value: Tensor<D>) -> Self {
130        // A plain backend can't activate gradients immediately, so record the setting explicitly
131        // for a later transition to training.
132        let mut param =
133            Param::initialized(ParamId::new(), Parameter::set_require_grad(value, true));
134        param.is_active = true;
135        param
136    }
137
138    /// Create a new parameter from data.
139    pub fn from_data<T>(data: T, device: &Device) -> Self
140    where
141        T: Into<TensorData>,
142    {
143        let data: TensorData = data.into();
144        // A plain backend can't activate gradients immediately, so record the setting explicitly
145        // for a later transition to training.
146        device.memory_persistent_allocations(data, |data| {
147            let value = Tensor::from_data(data, device);
148            let mut param =
149                Param::initialized(ParamId::new(), Parameter::set_require_grad(value, true));
150            param.is_active = true;
151            param
152        })
153    }
154
155    /// Attach a custom or built-in reparameterization, replacing any existing one.
156    pub(crate) fn with_reparameterization<R>(mut self, reparameterization: R) -> Self
157    where
158        R: Reparameterization,
159    {
160        self.reparameterization = Some(reparameterization_dyn::boxed::<R, D>(reparameterization));
161        self
162    }
163}
164
165impl<const D: usize> Module for Param<Tensor<D>> {
166    fn visit<V: ModuleVisitor>(&self, visitor: &mut V) {
167        match self.reparameterization_dyn() {
168            None => visitor.visit_float(self),
169            Some(reparameterization) => {
170                visitor.visit_float(&self.without_reparameterization());
171                visitor.enter_module(reparameterization.name(), "Reparameterization");
172                reparameterization_dyn::visit(reparameterization, visitor);
173                visitor.exit_module(reparameterization.name(), "Reparameterization");
174            }
175        }
176    }
177
178    fn map<M: ModuleMapper>(mut self, mapper: &mut M) -> Self {
179        match self.reparameterization.take() {
180            None => mapper.map_float(self),
181            Some(reparameterization) => {
182                let base = mapper.map_float(self);
183                mapper.enter_module(reparameterization.name(), "Reparameterization");
184                let reparameterization = reparameterization_dyn::map(reparameterization, mapper);
185                mapper.exit_module(reparameterization.name(), "Reparameterization");
186                base.with_dyn_reparameterization(Some(reparameterization))
187            }
188        }
189    }
190
191    fn to_device(mut self, device: &Device) -> Self {
192        let reparameterization = self.reparameterization.take();
193        let base = self.map_to_device(device, |tensor| tensor.to_device(device));
194        match reparameterization {
195            None => base,
196            Some(reparameterization) => {
197                base.with_dyn_reparameterization(Some(reparameterization.to_device_dyn(device)))
198            }
199        }
200    }
201
202    fn fork(mut self, device: &Device) -> Self {
203        let reparameterization = self.reparameterization.take();
204        let base = self.map_to_device(device, |tensor| {
205            let is_require_grad = tensor.is_require_grad();
206            let mut tensor = tensor.to_device(device).detach();
207
208            if is_require_grad {
209                tensor = tensor.require_grad();
210            }
211
212            tensor
213        });
214        match reparameterization {
215            None => base,
216            Some(reparameterization) => {
217                base.with_dyn_reparameterization(Some(reparameterization.fork_dyn(device)))
218            }
219        }
220    }
221
222    fn collect_devices(&self, mut devices: Vec<Device>) -> Vec<Device> {
223        let device = self.lazy_device();
224
225        if !devices.contains(&device) {
226            devices.push(device)
227        }
228
229        if let Some(reparameterization) = self.reparameterization_dyn() {
230            devices = reparameterization.collect_devices_dyn(devices);
231        }
232
233        devices
234    }
235
236    fn valid(&self) -> Self {
237        // Preserve whether the parameter was active, but reset the inner value's gradient state.
238        // Convert the stored base and reparameterization separately so validation never
239        // evaluates the effective weight or replaces a packed base with a dense one.
240        let is_active = self.is_active;
241        let mut param = Param::from_mapped_value(
242            self.id,
243            self.base().without_autodiff().set_require_grad(false),
244            self.param_mapper.clone(),
245        );
246        param.is_active = is_active;
247        param.reparameterization = self.reparameterization_dyn().map(|state| state.valid_dyn());
248        param
249    }
250
251    fn materialize(self) -> Self {
252        if self.reparameterization.is_none() {
253            return self;
254        }
255
256        // The merged weight is a new leaf with the base's training state. Keeping the adapter
257        // graph here would retain factors that are no longer exposed to the optimizer.
258        let require_grad = self.base().is_require_grad();
259        let value = self.val().detach().set_require_grad(require_grad);
260        let is_active = self.is_active;
261        let mut param = Param::from_mapped_value(self.id, value, self.param_mapper);
262        param.is_active = is_active;
263        param
264    }
265
266    fn train(mut self) -> Self {
267        // Keep the reparameterization structure and its parameters on the autodiff backend.
268        let reparameterization = self.reparameterization.take();
269        // Reinstate the parameter's training state.
270        let is_active = self.is_active;
271        let tensor = Tensor::from_inner(self.val()).set_require_grad(is_active);
272        let mut base = Param::from_mapped_value(self.id, tensor, self.param_mapper);
273        base.is_active = is_active;
274        match reparameterization {
275            None => base,
276            Some(reparameterization) => {
277                base.with_dyn_reparameterization(Some(reparameterization.train_dyn()))
278            }
279        }
280    }
281}
282
283impl<const D: usize> ModuleDisplayDefault for Param<Tensor<D>> {
284    fn content(&self, content: Content) -> Option<Content> {
285        let id = if content.display_settings.show_param_id() {
286            format!(", id: {}", self.id)
287        } else {
288            "".to_string()
289        };
290        let string = format!(
291            "ParamTensor {{rank: {D}, shape: {:?}, kind: float{id}}}",
292            self.shape().as_slice()
293        );
294        content.add_formatted(&string).optional()
295    }
296}
297impl<const D: usize> ModuleDisplay for Param<Tensor<D>> {}
298
299impl<const D: usize> Module for Param<Tensor<D, Int>> {
300    fn materialize(self) -> Self {
301        self
302    }
303
304    fn visit<V: ModuleVisitor>(&self, visitor: &mut V) {
305        visitor.visit_int(self)
306    }
307
308    fn map<M: ModuleMapper>(self, mapper: &mut M) -> Self {
309        mapper.map_int(self)
310    }
311
312    fn to_device(self, device: &Device) -> Self {
313        self.map_to_device(device, |tensor| tensor.to_device(device))
314    }
315
316    fn fork(self, device: &Device) -> Self {
317        self.to_device(device) // Don't support autodiff.
318    }
319
320    fn collect_devices(&self, mut devices: Vec<Device>) -> Vec<Device> {
321        let device = self.lazy_device();
322
323        if !devices.contains(&device) {
324            devices.push(device)
325        }
326
327        devices
328    }
329
330    fn valid(&self) -> Self {
331        Param::from_mapped_value(
332            self.id,
333            self.val().without_autodiff(),
334            self.param_mapper.clone(),
335        )
336    }
337
338    fn train(self) -> Self {
339        Param::from_mapped_value(self.id, Tensor::from_inner(self.val()), self.param_mapper)
340    }
341}
342
343impl<const D: usize> ModuleDisplayDefault for Param<Tensor<D, Int>> {
344    fn content(&self, content: Content) -> Option<Content> {
345        let id = if content.display_settings.show_param_id() {
346            format!(", id: {}", self.id)
347        } else {
348            "".to_string()
349        };
350        let string = format!(
351            "ParamTensor {{rank: {D}, shape: {:?}, kind: int{id}}}",
352            self.shape().as_slice()
353        );
354        content.add_formatted(&string).optional()
355    }
356}
357impl<const D: usize> ModuleDisplay for Param<Tensor<D, Int>> {}
358
359impl<const D: usize> Module for Param<Tensor<D, Bool>> {
360    fn materialize(self) -> Self {
361        self
362    }
363
364    fn visit<V: ModuleVisitor>(&self, visitor: &mut V) {
365        visitor.visit_bool(self)
366    }
367
368    fn map<M: ModuleMapper>(self, mapper: &mut M) -> Self {
369        mapper.map_bool(self)
370    }
371
372    fn to_device(self, device: &Device) -> Self {
373        self.map_to_device(device, |tensor| tensor.to_device(device))
374    }
375
376    fn fork(self, device: &Device) -> Self {
377        self.to_device(device) // Don't support autodiff.
378    }
379
380    fn collect_devices(&self, mut devices: Vec<Device>) -> Vec<Device> {
381        let device = self.lazy_device();
382
383        if !devices.contains(&device) {
384            devices.push(device)
385        }
386
387        devices
388    }
389
390    fn valid(&self) -> Self {
391        Param::from_mapped_value(
392            self.id,
393            self.val().without_autodiff(),
394            self.param_mapper.clone(),
395        )
396    }
397
398    fn train(self) -> Self {
399        Param::from_mapped_value(self.id, Tensor::from_inner(self.val()), self.param_mapper)
400    }
401}
402
403impl<const D: usize> ModuleDisplayDefault for Param<Tensor<D, Bool>> {
404    fn content(&self, content: Content) -> Option<Content> {
405        let id = if content.display_settings.show_param_id() {
406            format!(", id: {}", self.id)
407        } else {
408            "".to_string()
409        };
410
411        let string = format!(
412            "ParamTensor {{rank: {D}, shape: {:?}, kind: bool{id}}}",
413            self.shape().as_slice()
414        );
415        content.add_formatted(&string).optional()
416    }
417}
418
419impl<const D: usize> ModuleDisplay for Param<Tensor<D, Bool>> {}
420
421#[cfg(all(test, feature = "std", feature = "autodiff"))]
422mod tests {
423    use super::*;
424    use crate as burn;
425    use crate::{
426        module::{LoraAdapter, Module},
427        test_device,
428    };
429    use burn_tensor::Distribution;
430
431    #[test]
432    fn validation_does_not_apply_reparameterizations() {
433        #[derive(Debug, Module)]
434        struct NeverApply {}
435
436        impl Reparameterization for NeverApply {
437            const NAME: &'static str = "never_apply";
438
439            fn apply<const D: usize>(&self, _base: Tensor<D>) -> Tensor<D> {
440                panic!("validation must not evaluate effective weights")
441            }
442        }
443
444        let device = test_device().autodiff();
445        let param = Param::from_tensor(Tensor::<2>::ones([2, 3], &device))
446            .with_reparameterization(NeverApply {});
447        let validation = param.valid().valid();
448        assert!(!validation.base().is_autodiff());
449        assert!(validation.reparameterization::<NeverApply>().is_some());
450    }
451
452    #[test]
453    fn materialize_detaches_factors_and_preserves_base_state_and_layout_mapping() {
454        let device = test_device().autodiff();
455        let a = Param::from_tensor(Tensor::<2>::ones([2, 1], &device));
456        let b = Param::from_tensor(Tensor::<2>::ones([1, 3], &device));
457        let param = Param::from_tensor(Tensor::<2>::ones([2, 3], &device))
458            .load_mapper(Tensor::transpose)
459            .save_mapper(Tensor::transpose)
460            .with_reparameterization(LoraAdapter {
461                a: a.clone(),
462                b: b.clone(),
463                scale: 2.0,
464            });
465
466        let merged = param.clone().materialize();
467        assert!(merged.adapter().is_none());
468        assert!(param.adapter().is_some());
469        assert_eq!(merged.id, param.id);
470        assert!(merged.base().is_require_grad());
471        let value = merged.base().into_data();
472        value.assert_eq(&TensorData::from([[3.0f32; 3]; 2]), true);
473        let saved = merged.transform_for_save().base().into_data();
474        saved.assert_eq(&TensorData::from([[3.0f32; 2]; 3]), true);
475        let loaded = merged
476            .clone()
477            .transform_for_load(Tensor::ones([3, 2], &device), merged.id);
478        assert_eq!(loaded.base().dims(), [2, 3]);
479
480        let grads = merged.val().sum().backward();
481        assert!(merged.base().grad(&grads).is_some());
482        assert!(param.base().grad(&grads).is_none());
483        assert!(a.val().grad(&grads).is_none());
484        assert!(b.val().grad(&grads).is_none());
485
486        let validation_merged = param.valid().materialize();
487        assert!(!validation_merged.base().is_autodiff());
488        assert!(validation_merged.train().base().is_require_grad());
489        let frozen_merged = param.freeze().materialize();
490        assert!(!frozen_merged.base().is_require_grad());
491        assert!(!frozen_merged.valid().train().base().is_require_grad());
492    }
493
494    #[test]
495    fn qlora_validation_and_records_preserve_packed_base_until_materialization() {
496        use crate::module::{Lora, QLora, Quantizer};
497        use burn_tensor::quantization::{Calibration, QuantValue};
498
499        let device = test_device().autodiff();
500        let scheme = device
501            .settings()
502            .quantization
503            .scheme
504            .with_value(QuantValue::Q8S);
505        let param = Param::from_tensor(Tensor::<2>::ones([4, 4], &device)).apply_qlora(QLora::new(
506            Lora::new(2, 4.0),
507            Quantizer::new(Calibration::MinMax, scheme),
508        ));
509        let inference = param.valid();
510        assert_eq!(inference.base().dtype(), param.base().dtype());
511        let dtype = inference.base().dtype();
512        assert!(matches!(dtype, burn_tensor::DType::QFloat(_)));
513        assert!(inference.adapter().is_some());
514        let restored = inference.clone().train();
515        assert!(!restored.base().is_require_grad());
516        let adapter = restored.adapter().unwrap();
517        assert!(adapter.a.val().is_require_grad());
518        assert!(adapter.b.val().is_require_grad());
519
520        let expected = inference.val().into_data();
521        let loaded = param.load_record(inference.clone().into_record());
522        assert_eq!(loaded.base().dtype(), inference.base().dtype());
523        assert!(loaded.adapter().is_some());
524        loaded.val().into_data().assert_eq(&expected, true);
525
526        let merged = inference.materialize().materialize();
527        assert!(merged.adapter().is_none());
528        assert!(merged.base().dtype().is_float());
529        assert!(!merged.base().is_autodiff());
530        merged.val().into_data().assert_eq(&expected, true);
531        let dense = Param::from_tensor(Tensor::<2>::zeros([4, 4], &test_device()));
532        let loaded = dense.load_record(merged.into_record());
533        loaded.val().into_data().assert_eq(&expected, true);
534    }
535
536    fn lazy_param_for_device_inspection<T: Parameter>(
537        device: &Device,
538        shape: [usize; 2],
539    ) -> Param<T> {
540        Param::uninitialized(
541            ParamId::new(),
542            |_, _| panic!("device inspection must not initialize parameters"),
543            device.clone(),
544            false,
545            shape.into(),
546        )
547    }
548
549    #[test]
550    fn devices_preserves_lazy_parameters_of_all_tensor_kinds() {
551        #[derive(Module, Debug)]
552        struct MixedParameters {
553            float: Param<Tensor<2>>,
554            int: Param<Tensor<2, Int>>,
555            bool: Param<Tensor<2, Bool>>,
556            initialized: Param<Tensor<2>>,
557        }
558
559        let device = test_device();
560        let model = MixedParameters {
561            float: lazy_param_for_device_inspection(&device, [2, 3]),
562            int: lazy_param_for_device_inspection(&device, [2, 3]),
563            bool: lazy_param_for_device_inspection(&device, [2, 3]),
564            initialized: Param::from_tensor(Tensor::ones([2, 3], &device)),
565        };
566        let clone = model.clone();
567
568        assert_eq!(model.devices(), alloc::vec![device.clone()]);
569        assert_eq!(clone.devices(), alloc::vec![device]);
570        assert!(!model.float.is_initialized());
571        assert!(!model.int.is_initialized());
572        assert!(!model.bool.is_initialized());
573    }
574
575    #[test]
576    fn devices_preserves_lazy_reparameterization_state() {
577        let device = test_device();
578        let param: Param<Tensor<2>> = lazy_param_for_device_inspection(&device, [2, 3]);
579
580        // Cover both lazy and initialized bases with lazy adapter parameters.
581        for base in [
582            param,
583            Param::from_tensor(Tensor::<2>::ones([2, 3], &device)),
584        ] {
585            let was_initialized = base.is_initialized();
586            let param = base.with_reparameterization(LoraAdapter {
587                a: lazy_param_for_device_inspection(&device, [2, 1]),
588                b: lazy_param_for_device_inspection(&device, [1, 3]),
589                scale: 1.0,
590            });
591
592            assert_eq!(param.devices(), alloc::vec![device.clone()]);
593            assert_eq!(param.is_initialized(), was_initialized);
594            let adapter = param.adapter().unwrap();
595            assert!(!adapter.a.is_initialized());
596            assert!(!adapter.b.is_initialized());
597        }
598    }
599
600    #[test]
601    fn devices_preserves_gradients_after_a_lazy_module_move() {
602        let device = test_device().autodiff();
603        let model = crate::test_utils::SimpleLinear {
604            weight: Param::uninitialized(
605                ParamId::new(),
606                |device, require_grad| Tensor::ones([1, 2], device).set_require_grad(require_grad),
607                device.clone(),
608                true,
609                [1, 2].into(),
610            ),
611            bias: Some(Param::uninitialized(
612                ParamId::new(),
613                |device, require_grad| Tensor::zeros([1], device).set_require_grad(require_grad),
614                device.clone(),
615                true,
616                [1].into(),
617            )),
618        };
619
620        assert_eq!(model.devices(), alloc::vec![device.clone()]);
621        assert!(!model.weight.is_initialized());
622        assert!(!model.bias.as_ref().unwrap().is_initialized());
623
624        let model = model.to_device(&device);
625        let output = Tensor::<2>::ones([1, 2], &device).matmul(model.weight.val().transpose())
626            + model.bias.as_ref().unwrap().val().unsqueeze();
627        let grads = output.sum().backward();
628
629        model
630            .weight
631            .grad(&grads)
632            .unwrap()
633            .into_data()
634            .assert_eq(&TensorData::from([[1.0f32, 1.0f32]]), false);
635        model
636            .bias
637            .as_ref()
638            .unwrap()
639            .grad(&grads)
640            .unwrap()
641            .into_data()
642            .assert_eq(&TensorData::from([1.0f32]), false);
643    }
644
645    #[test]
646    fn set_require_grad_updates_lazy_lifecycle_state() {
647        let device = test_device().autodiff();
648        let param: Param<Tensor<2>> = Param::uninitialized(
649            ParamId::new(),
650            |device, require_grad| Tensor::ones([2, 3], device).set_require_grad(require_grad),
651            device,
652            true,
653            [2, 3].into(),
654        );
655
656        let param = param.set_require_grad(false);
657
658        assert!(!param.is_initialized());
659        assert!(!param.is_active);
660        assert!(!param.val().is_require_grad());
661
662        let param = param.valid().train();
663
664        assert!(!param.is_require_grad());
665        assert!(!param.is_active);
666    }
667
668    #[test]
669    fn set_require_grad_on_a_plain_tensor_is_applied_by_train() {
670        let device = test_device();
671        let param = Param::initialized(
672            ParamId::new(),
673            Tensor::<2>::ones([2, 3], &device).set_require_grad(false),
674        )
675        .set_require_grad(true);
676
677        assert!(!param.is_require_grad());
678        assert!(param.is_active);
679
680        let param = param.train();
681
682        assert!(param.is_require_grad());
683        assert!(param.is_active);
684    }
685
686    #[test]
687    fn trainable_param_created_on_a_plain_device_is_applied_by_train() {
688        let device = test_device();
689        let param = Param::from_tensor(Tensor::<2>::ones([2, 3], &device));
690
691        assert!(!param.is_require_grad());
692        assert!(param.is_active);
693
694        let param = param.train();
695
696        assert!(param.is_require_grad());
697        assert!(param.is_active);
698    }
699
700    #[test]
701    fn lazy_activation_setting_on_a_plain_device_is_applied_by_train() {
702        let device = test_device();
703        let param: Param<Tensor<2>> = Param::uninitialized(
704            ParamId::new(),
705            |device, require_grad| Tensor::ones([2, 3], device).set_require_grad(require_grad),
706            device,
707            false,
708            [2, 3].into(),
709        )
710        .set_require_grad(true);
711
712        assert!(!param.is_initialized());
713        assert!(param.is_active);
714        assert!(!param.val().is_require_grad());
715
716        let param = param.train();
717
718        assert!(param.is_require_grad());
719        assert!(param.is_active);
720    }
721
722    #[test]
723    fn mapping_a_validation_param_preserves_what_train_restores() {
724        let device = test_device().autodiff();
725        let param = Param::from_tensor(Tensor::<2>::ones([2, 3], &device))
726            .valid()
727            .map(|tensor| tensor);
728
729        assert!(!param.is_require_grad());
730        assert!(param.is_active);
731
732        let param = param.train();
733
734        assert!(param.is_require_grad());
735        assert!(param.is_active);
736    }
737
738    #[test]
739    fn mapped_value_reconstruction_preserves_what_train_restores() {
740        let device = test_device().autodiff();
741        let valid = Param::from_tensor(Tensor::<2>::ones([2, 3], &device)).valid();
742        let (id, tensor, mapper) = valid.consume();
743
744        let param = Param::from_mapped_value(id, tensor, mapper);
745
746        assert!(!param.is_require_grad());
747        assert!(param.is_active);
748
749        let param = param.train();
750
751        assert!(param.is_require_grad());
752        assert!(param.is_active);
753    }
754
755    #[test]
756    fn loading_a_validation_param_preserves_what_train_restores() {
757        let device = test_device().autodiff();
758        let valid = Param::from_tensor(Tensor::<2>::ones([2, 3], &device)).valid();
759        let record = Tensor::<2>::zeros([2, 3], &test_device());
760
761        let param = valid.transform_for_load(record, ParamId::new());
762
763        assert!(!param.is_require_grad());
764        assert!(param.is_active);
765
766        let param = param.train();
767
768        assert!(param.is_require_grad());
769        assert!(param.is_active);
770    }
771
772    #[test]
773    fn test_param_require_grad_stateful() {
774        let device = test_device().autodiff();
775        let tensor = Tensor::<2>::ones([3, 3], &device).require_grad();
776
777        let param = Param::initialized(ParamId::new(), tensor);
778        assert!(param.is_require_grad());
779        assert!(param.is_active);
780
781        let param = param.valid();
782        assert!(!param.is_require_grad());
783        assert!(param.is_active); // stateful
784
785        let param = param.train();
786        assert!(param.is_require_grad());
787        assert!(param.is_active); // stateful
788
789        let param = param.no_grad();
790        assert!(!param.is_require_grad());
791        assert!(!param.is_active); // stateful
792
793        let param = param.valid();
794        assert!(!param.is_require_grad()); // always
795        assert!(!param.is_active); // stateful
796
797        let param = param.train();
798        assert!(!param.is_require_grad());
799        assert!(!param.is_active); // stateful
800    }
801
802    #[test]
803    fn a_lazy_param_with_an_init_mapper_trains_on_an_autodiff_device() {
804        let device = test_device().autodiff();
805        let param: Param<Tensor<2>> = Param::uninitialized(
806            ParamId::new(),
807            |device, require_grad| Tensor::ones([2, 3], device).set_require_grad(require_grad),
808            device,
809            true,
810            [2, 3].into(),
811        )
812        .init_mapper(|tensor| tensor.mul_scalar(2.0));
813
814        let value = param.val();
815        let grads = value.clone().sum().backward();
816
817        value
818            .into_data()
819            .assert_eq(&TensorData::from([[2.0f32; 3]; 2]), false);
820        param
821            .grad(&grads)
822            .expect("the mapped value is the leaf that receives the gradient")
823            .into_data()
824            .assert_eq(&TensorData::from([[1.0f32; 3]; 2]), false);
825    }
826
827    #[test]
828    fn counting_and_materializing_an_unadapted_lazy_param_leave_it_uninitialized() {
829        let param: Param<Tensor<2>> = Param::uninitialized(
830            ParamId::new(),
831            |device, require_grad| Tensor::ones([2, 3], device).set_require_grad(require_grad),
832            test_device(),
833            true,
834            [2, 3].into(),
835        );
836
837        assert_eq!(Module::num_params(&param), 6);
838        assert!(!param.is_initialized());
839        let id = param.id;
840        let param = param.materialize().materialize();
841        assert_eq!(param.id, id);
842        assert!(!param.is_initialized());
843    }
844
845    #[test]
846    fn a_lazy_param_moved_then_loaded_never_initializes() {
847        let device = test_device();
848        let param: Param<Tensor<2>> = Param::uninitialized(
849            ParamId::new(),
850            |_, _| panic!("the moved parameter initialized before loading"),
851            device.clone(),
852            false,
853            [2, 3].into(),
854        );
855
856        let moved = param.to_device(&device.clone().autodiff());
857
858        assert!(!moved.is_initialized());
859        moved.transform_for_load(Tensor::ones([2, 3], &device), ParamId::new());
860    }
861
862    #[test]
863    fn a_moved_lazy_param_keeps_the_autodiff_context_it_was_built_with() {
864        let device = test_device();
865        let lazy_ones = |device: &Device| -> Param<Tensor<2>> {
866            Param::uninitialized(
867                ParamId::new(),
868                |device, require_grad| Tensor::ones([2, 3], device).set_require_grad(require_grad),
869                device.clone(),
870                true,
871                [2, 3].into(),
872            )
873        };
874
875        let onto_autodiff = lazy_ones(&device).to_device(&device.clone().autodiff());
876        let onto_plain = lazy_ones(&device.clone().autodiff()).fork(&device);
877
878        assert!(!onto_autodiff.lazy_device().is_autodiff());
879        assert!(onto_plain.lazy_device().is_autodiff());
880    }
881
882    #[test]
883    fn a_moved_param_keeps_its_reparameterization() {
884        let device = test_device();
885        let adapter = || LoraAdapter {
886            a: Param::from_tensor(Tensor::<2>::ones([3, 1], &device)),
887            b: Param::from_tensor(Tensor::<2>::zeros([1, 3], &device)),
888            scale: 1.0,
889        };
890        let lazy: Param<Tensor<2>> = Param::uninitialized(
891            ParamId::new(),
892            |device, _| Tensor::ones([3, 3], device),
893            device.clone(),
894            false,
895            [3, 3].into(),
896        )
897        .with_reparameterization(adapter());
898        let initialized = Param::from_tensor(Tensor::<2>::ones([3, 3], &device))
899            .with_reparameterization(adapter());
900
901        let target = device.clone().autodiff();
902
903        assert!(lazy.to_device(&target).adapter().is_some());
904        assert!(initialized.fork(&target).adapter().is_some());
905    }
906
907    #[test]
908    fn init_mapper_preserves_lora_after_a_clone_initializes_the_base() {
909        let device = test_device();
910        let param: Param<Tensor<2>> = Param::uninitialized(
911            ParamId::new(),
912            |device, _| Tensor::ones([3, 3], device),
913            device.clone(),
914            false,
915            [3, 3].into(),
916        )
917        .with_reparameterization(crate::module::LoraAdapter {
918            a: Param::from_tensor(Tensor::<2>::ones([3, 1], &device)),
919            b: Param::from_tensor(Tensor::<2>::ones([1, 3], &device)),
920            scale: 2.0,
921        });
922        let clone = param.clone();
923        assert!(!param.is_initialized());
924        let mapped = param.init_mapper(|value| value.mul_scalar(2.0));
925
926        // Initialize the captured base after registering the lazy mapper, then remove sharing.
927        // This exercises init_mapper's map_to_device fallback, not Module::to_device or fork.
928        clone
929            .val()
930            .into_data()
931            .assert_eq(&TensorData::from([[3.0f32; 3]; 3]), true);
932        drop(clone);
933        assert!(!mapped.is_initialized());
934
935        // Map the effective LoRA value: (base + scale * A @ B) * 2 = (1 + 2) * 2 = 6.
936        mapped
937            .val()
938            .into_data()
939            .assert_eq(&TensorData::from([[6.0f32; 3]; 3]), true);
940    }
941
942    #[test]
943    fn a_lazy_int_param_moved_never_initializes() {
944        let device = test_device();
945        let param: Param<Tensor<2, Int>> = Param::uninitialized(
946            ParamId::new(),
947            |_, _| panic!("the moved parameter initialized"),
948            device.clone(),
949            false,
950            [2, 3].into(),
951        );
952
953        let param = param.to_device(&device.clone().autodiff());
954
955        assert!(!param.is_initialized());
956        assert!(!param.lazy_device().is_autodiff());
957    }
958
959    #[test]
960    fn a_lazy_bool_param_moved_never_initializes() {
961        let device = test_device();
962        let param: Param<Tensor<2, Bool>> = Param::uninitialized(
963            ParamId::new(),
964            |_, _| panic!("the moved parameter initialized"),
965            device.clone(),
966            false,
967            [2, 3].into(),
968        );
969
970        let param = param.to_device(&device.clone().autodiff());
971
972        assert!(!param.is_initialized());
973        assert!(!param.lazy_device().is_autodiff());
974    }
975
976    #[test]
977    fn a_lazy_param_forked_keeps_its_gradient_requirement() {
978        let device = test_device();
979        let param: Param<Tensor<2>> = Param::uninitialized(
980            ParamId::new(),
981            |device, require_grad| Tensor::ones([2, 3], device).set_require_grad(require_grad),
982            device.clone().autodiff(),
983            true,
984            [2, 3].into(),
985        );
986
987        let param = param.fork(&device);
988
989        assert!(!param.is_initialized());
990        assert!(param.val().is_require_grad());
991    }
992
993    #[test]
994    fn a_lazy_param_with_an_init_mapper_follows_the_move() {
995        let device = test_device().autodiff();
996        let param: Param<Tensor<2>> = Param::uninitialized(
997            ParamId::new(),
998            |device, require_grad| Tensor::ones([2, 3], device).set_require_grad(require_grad),
999            device.clone(),
1000            true,
1001            [2, 3].into(),
1002        )
1003        .init_mapper(|tensor| tensor.mul_scalar(2.0));
1004
1005        let param = param.fork(&device);
1006
1007        assert!(!param.is_initialized());
1008        let value = param.val();
1009        assert!(value.device().is_autodiff());
1010        assert!(value.is_require_grad());
1011    }
1012
1013    #[test]
1014    fn a_lazy_param_shared_with_a_clone_initializes_before_moving() {
1015        let device = test_device();
1016        let param: Param<Tensor<2>> = Param::uninitialized(
1017            ParamId::new(),
1018            |device, _| Tensor::random([2, 3], Distribution::Default, device),
1019            device.clone(),
1020            false,
1021            [2, 3].into(),
1022        );
1023        let clone = param.clone();
1024
1025        let moved = param.to_device(&device.autodiff());
1026
1027        assert!(clone.is_initialized());
1028        moved
1029            .val()
1030            .into_data()
1031            .assert_eq(&clone.val().into_data(), true);
1032    }
1033}
burn_core/module/param/tensor.rs

burn_core/module/param/
tensor.rs
